mirror of
https://github.com/Refound-445/nonebot-plugin-nailongremove.git
synced 2025-09-27 02:49:17 +08:00
chore: pre-commit auto fix [skip ci]
This commit is contained in:
parent
8a55fed055
commit
6fb53ca9e2
@ -105,7 +105,7 @@ async def handle_function(bot: BaseBot, ev: BaseEvent, msg: UniMsg, session: Uni
|
||||
frames.append(temp_image)
|
||||
except StopIteration:
|
||||
break
|
||||
commitInfo = process_gif_and_save_jpgs(frames, label, (224,224))
|
||||
commitInfo = process_gif_and_save_jpgs(frames, label, (224, 224))
|
||||
if commitInfo is None:
|
||||
await nailong.finish(
|
||||
f"The new data has been saved to the directory {config.nailong_model_dir}\\records\\{label}, label: {label}.",
|
||||
@ -137,7 +137,9 @@ async def handle_function(bot: BaseBot, ev: BaseEvent, msg: UniMsg, session: Uni
|
||||
]
|
||||
if len(template_str_all) == 0:
|
||||
continue
|
||||
template_str=template_str_all[random.randint(0, len(template_str_all) - 1)]
|
||||
template_str = template_str_all[
|
||||
random.randint(0, len(template_str_all) - 1)
|
||||
]
|
||||
mapping = {
|
||||
"$event": ev,
|
||||
"$target": msg.get_target(),
|
||||
|
@ -38,7 +38,6 @@ else:
|
||||
file_path = os.path.join(str(config.nailong_model_dir), FILENAME)
|
||||
model_info = api.model_info(REPO_ID)
|
||||
|
||||
|
||||
def get_file_last_modified_time(file_path):
|
||||
try:
|
||||
timestamp = os.path.getmtime(file_path)
|
||||
@ -50,7 +49,6 @@ else:
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
|
||||
|
||||
local_time = get_file_last_modified_time(file_path)
|
||||
if local_time is None or model_info.last_modified >= local_time:
|
||||
hf_hub_download(
|
||||
@ -78,7 +76,7 @@ def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult:
|
||||
input_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
|
||||
if not os.path.exists(
|
||||
os.path.join(str(config.nailong_model_dir), "online_temp"),
|
||||
os.path.join(str(config.nailong_model_dir), "online_temp"),
|
||||
):
|
||||
os.makedirs(os.path.join(str(config.nailong_model_dir), "online_temp"))
|
||||
image_path = os.path.join(
|
||||
@ -102,8 +100,8 @@ def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult:
|
||||
)
|
||||
os.remove(image_path)
|
||||
if (
|
||||
"检测到的目标数量: " in result_info
|
||||
and int(result_info.split("检测到的目标数量: ")[1].split("\n")[0]) < 1
|
||||
"检测到的目标数量: " in result_info
|
||||
and int(result_info.split("检测到的目标数量: ")[1].split("\n")[0]) < 1
|
||||
):
|
||||
return CheckSingleResult(ok=False, label=None, extra=frame)
|
||||
if isinstance(result_image, str):
|
||||
@ -144,9 +142,9 @@ def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult:
|
||||
|
||||
if pad_w > 0 or pad_h > 0:
|
||||
result_img = result_img[
|
||||
pad_h // 2: pad_h // 2 + original_size[1],
|
||||
pad_w // 2: pad_w // 2 + original_size[0],
|
||||
]
|
||||
pad_h // 2 : pad_h // 2 + original_size[1],
|
||||
pad_w // 2 : pad_w // 2 + original_size[0],
|
||||
]
|
||||
return CheckSingleResult(ok=True, label="nailong", extra=result_img)
|
||||
|
||||
|
||||
|
@ -83,8 +83,8 @@ class FrameInfo:
|
||||
|
||||
@run_sync
|
||||
def _check_single(
|
||||
frame: np.ndarray,
|
||||
is_gif: bool = False,
|
||||
frame: np.ndarray,
|
||||
is_gif: bool = False,
|
||||
) -> CheckSingleResult[Optional[Detections]]:
|
||||
if is_gif:
|
||||
res = similarity_process(frame)
|
||||
@ -127,8 +127,8 @@ def _check_single(
|
||||
|
||||
|
||||
async def check_single(
|
||||
frame: np.ndarray,
|
||||
is_gif: bool = False,
|
||||
frame: np.ndarray,
|
||||
is_gif: bool = False,
|
||||
) -> CheckSingleResult[FrameInfo]:
|
||||
if is_gif:
|
||||
res = await _check_single(frame, True)
|
||||
|
@ -7,13 +7,12 @@ import random
|
||||
import shutil
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Awaitable, Callable, Dict, Generic, Optional, TypeVar
|
||||
|
||||
from typing_extensions import TypeAlias
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ...config import config
|
||||
from ...frame_source import FrameSource
|
||||
|
||||
@ -22,20 +21,24 @@ T = TypeVar("T")
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
if config.nailong_similarity_on:
|
||||
from huggingface_hub import PyTorchModelHubMixin
|
||||
from torch import nn
|
||||
import torchvision
|
||||
from nonebot import logger
|
||||
import faiss
|
||||
import json
|
||||
import sklearn
|
||||
|
||||
import faiss
|
||||
import torchvision
|
||||
from huggingface_hub import PyTorchModelHubMixin
|
||||
from nonebot import logger
|
||||
from torch import nn
|
||||
from torchvision import transforms
|
||||
|
||||
transform = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.5], std=[0.5]) # Assuming grayscale or single-channel
|
||||
])
|
||||
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(
|
||||
mean=[0.5],
|
||||
std=[0.5],
|
||||
), # Assuming grayscale or single-channel
|
||||
],
|
||||
)
|
||||
|
||||
class MyModel(
|
||||
nn.Module,
|
||||
@ -44,21 +47,25 @@ if config.nailong_similarity_on:
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.resnet = torchvision.models.resnet18(pretrained=False)
|
||||
self.resnet.fc = nn.Linear(self.resnet.fc.in_features, 5) # Output dimension is 5
|
||||
self.resnet.fc = nn.Linear(
|
||||
self.resnet.fc.in_features,
|
||||
5,
|
||||
) # Output dimension is 5
|
||||
|
||||
def forward(self, x):
|
||||
return self.resnet(x)
|
||||
|
||||
|
||||
features_model = MyModel.from_pretrained("refoundd/NailongFeatures", ).to(device)
|
||||
index_path = config.nailong_model_dir / 'records.index'
|
||||
json_path = config.nailong_model_dir / 'records.json'
|
||||
features_model = MyModel.from_pretrained(
|
||||
"refoundd/NailongFeatures",
|
||||
).to(device)
|
||||
index_path = config.nailong_model_dir / "records.index"
|
||||
json_path = config.nailong_model_dir / "records.json"
|
||||
if os.path.exists(index_path):
|
||||
index = faiss.read_index(str(index_path))
|
||||
else:
|
||||
index = faiss.IndexFlatL2(512)
|
||||
if os.path.exists(json_path):
|
||||
with open(json_path, 'r') as f:
|
||||
with open(json_path, "r") as f:
|
||||
index_cls = json.load(f)
|
||||
else:
|
||||
index_cls = {}
|
||||
@ -66,9 +73,10 @@ if config.nailong_similarity_on:
|
||||
try:
|
||||
res = faiss.StandardGpuResources() # 创建GPU资源
|
||||
index = faiss.index_cpu_to_gpu(res, 0, index) # 将CPU索引转移到GPU
|
||||
except Exception as e:
|
||||
logger.warning("load faiss-gpu failed.Please check your GPU device and install faiss-gpu first.")
|
||||
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"load faiss-gpu failed.Please check your GPU device and install faiss-gpu first.",
|
||||
)
|
||||
|
||||
def hook(model, input, output):
|
||||
embeddings = input[0]
|
||||
@ -78,7 +86,6 @@ if config.nailong_similarity_on:
|
||||
d, i = index.search(vector, 1)
|
||||
return 1 - d[0][0], i[0][0], vector
|
||||
|
||||
|
||||
features_model.resnet.fc.register_forward_hook(hook)
|
||||
features_model.eval()
|
||||
|
||||
@ -112,9 +119,9 @@ FrameChecker: TypeAlias = Callable[
|
||||
|
||||
|
||||
async def race_check(
|
||||
checker: FrameChecker[T],
|
||||
frames: FrameSource,
|
||||
concurrency: int = config.nailong_concurrency,
|
||||
checker: FrameChecker[T],
|
||||
frames: FrameSource,
|
||||
concurrency: int = config.nailong_concurrency,
|
||||
) -> Optional[CheckSingleResult[T]]:
|
||||
iterator = iter(frames)
|
||||
if config.nailong_similarity_on:
|
||||
@ -182,7 +189,11 @@ async def race_check(
|
||||
return None
|
||||
|
||||
|
||||
def similarity_process(image1: np.ndarray, dsize=(224, 224), similarity_threshold=1) -> Optional[CheckSingleResult]:
|
||||
def similarity_process(
|
||||
image1: np.ndarray,
|
||||
dsize=(224, 224),
|
||||
similarity_threshold=1,
|
||||
) -> Optional[CheckSingleResult]:
|
||||
# image1 = cv2.cvtColor(image1, cv2.COLOR_BGR2RGB)
|
||||
image1 = cv2.resize(image1, dsize, interpolation=cv2.INTER_LINEAR)
|
||||
image1_tensor = transform(image1).unsqueeze(0).to(device)
|
||||
@ -198,22 +209,25 @@ def similarity_process(image1: np.ndarray, dsize=(224, 224), similarity_threshol
|
||||
|
||||
def process_gif_and_save_jpgs(frames, label, dsize=(224, 224), similarity_threshold=1):
|
||||
if (
|
||||
len(
|
||||
list(
|
||||
glob.glob(
|
||||
str(config.nailong_model_dir / "records/*/*.jpg")
|
||||
),
|
||||
),
|
||||
)
|
||||
>= config.nailong_similarity_max_storage and config.nailong_hf_token is not None
|
||||
len(
|
||||
list(
|
||||
glob.glob(str(config.nailong_model_dir / "records/*/*.jpg")),
|
||||
),
|
||||
)
|
||||
>= config.nailong_similarity_max_storage
|
||||
and config.nailong_hf_token is not None
|
||||
):
|
||||
zip_filename = shutil.make_archive(
|
||||
config.nailong_model_dir / "{}_records".format(datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")),
|
||||
config.nailong_model_dir
|
||||
/ "{}_records".format(
|
||||
datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
|
||||
),
|
||||
"zip",
|
||||
config.nailong_model_dir / "records"
|
||||
config.nailong_model_dir / "records",
|
||||
)
|
||||
shutil.rmtree(config.nailong_model_dir / "records")
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
api = HfApi()
|
||||
commitInfo = api.upload_file(
|
||||
path_or_fileobj=zip_filename,
|
||||
@ -255,6 +269,6 @@ def process_gif_and_save_jpgs(frames, label, dsize=(224, 224), similarity_thresh
|
||||
index_cls[str(index.ntotal - 1)] = label
|
||||
count += 1
|
||||
faiss.write_index(index, str(index_path))
|
||||
with open(json_path, 'w') as f:
|
||||
with open(json_path, "w") as f:
|
||||
json.dump(index_cls, f)
|
||||
return commitInfo
|
||||
|
Loading…
Reference in New Issue
Block a user