diff --git a/.idea/.gitignore b/.idea/.gitignore
new file mode 100644
index 0000000..35410ca
--- /dev/null
+++ b/.idea/.gitignore
@@ -0,0 +1,8 @@
+# 默认忽略的文件
+/shelf/
+/workspace.xml
+# 基于编辑器的 HTTP 客户端请求
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml
diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml
new file mode 100644
index 0000000..105ce2d
--- /dev/null
+++ b/.idea/inspectionProfiles/profiles_settings.xml
@@ -0,0 +1,6 @@
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.idea/misc.xml b/.idea/misc.xml
new file mode 100644
index 0000000..060d2c5
--- /dev/null
+++ b/.idea/misc.xml
@@ -0,0 +1,4 @@
+
+
+
+
\ No newline at end of file
diff --git a/.idea/modules.xml b/.idea/modules.xml
new file mode 100644
index 0000000..2e920b1
--- /dev/null
+++ b/.idea/modules.xml
@@ -0,0 +1,8 @@
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.idea/nonebot-plugin-nailongremove.iml b/.idea/nonebot-plugin-nailongremove.iml
new file mode 100644
index 0000000..d0876a7
--- /dev/null
+++ b/.idea/nonebot-plugin-nailongremove.iml
@@ -0,0 +1,8 @@
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
new file mode 100644
index 0000000..35eb1dd
--- /dev/null
+++ b/.idea/vcs.xml
@@ -0,0 +1,6 @@
+
+
+
+
+
+
\ No newline at end of file
diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/hf_detection.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/hf_detection.py
index 0e95ae0..e730aa1 100644
--- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/hf_detection.py
+++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/hf_detection.py
@@ -38,6 +38,7 @@ 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)
@@ -49,6 +50,7 @@ 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(
@@ -76,7 +78,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(
@@ -100,8 +102,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):
@@ -142,9 +144,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)
diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/target_detection.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/target_detection.py
index ae4f92f..85d2724 100644
--- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/target_detection.py
+++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/target_detection.py
@@ -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)
diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/common.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/common.py
index ea4a3c8..2dc978b 100644
--- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/common.py
+++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/common.py
@@ -35,6 +35,8 @@ if config.nailong_similarity_on:
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5]) # Assuming grayscale or single-channel
])
+
+
class MyModel(
nn.Module,
PyTorchModelHubMixin,
@@ -46,36 +48,41 @@ if config.nailong_similarity_on:
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))
+ index = faiss.read_index(str(index_path))
else:
- index=faiss.IndexFlatL2(512)
+ index = faiss.IndexFlatL2(512)
if os.path.exists(json_path):
with open(json_path, 'r') as f:
- index_cls=json.load(f)
+ index_cls = json.load(f)
else:
- index_cls= {}
+ index_cls = {}
if torch.cuda.is_available():
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.")
+
+
def hook(model, input, output):
- embeddings=input[0]
+ embeddings = input[0]
vector = embeddings.detach().cpu().numpy().astype(np.float32)
faiss.normalize_L2(vector)
global index
d, i = index.search(vector, 1)
- return 1-d[0][0],i[0][0],vector
+ return 1 - d[0][0], i[0][0], vector
+
+
features_model.resnet.fc.register_forward_hook(hook)
features_model.eval()
-
@dataclass
class CheckSingleResult(Generic[T]):
ok: bool
@@ -105,9 +112,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:
@@ -175,40 +182,40 @@ 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)
# image1_tensor = (
# torch.tensor(image1, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)
# ).to(device)
- distance,indice,_=features_model(image1_tensor)
+ distance, indice, _ = features_model(image1_tensor)
if distance >= similarity_threshold:
- label =index_cls[str(indice)]
+ label = index_cls[str(indice)]
return CheckSingleResult(ok=True, label=label, extra=None)
return None
-def process_gif_and_save_jpgs(frames, label, dsize=(224,224), similarity_threshold=1):
+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")
+ len(
+ list(
+ glob.glob(
+ str(config.nailong_model_dir / "records/*/*.jpg")
+ ),
),
- ),
- )
- >= config.nailong_similarity_max_storage and config.nailong_hf_token is not None
+ )
+ >= 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"
)
shutil.rmtree(config.nailong_model_dir / "records")
from huggingface_hub import HfApi
api = HfApi()
- commitInfo=api.upload_file(
+ commitInfo = api.upload_file(
path_or_fileobj=zip_filename,
path_in_repo="new_dataset.zip",
repo_id="refoundd/NailongClassification",
@@ -219,10 +226,10 @@ def process_gif_and_save_jpgs(frames, label, dsize=(224,224), similarity_thresho
# os.remove(zip_filename)
else:
commitInfo = None
- output_dir = config.nailong_model_dir / "records"/ label
+ output_dir = config.nailong_model_dir / "records" / label
if not os.path.exists(output_dir):
os.makedirs(output_dir)
- count=0
+ count = 0
for frame in frames:
frame_filename = os.path.join(
output_dir,
@@ -235,19 +242,19 @@ def process_gif_and_save_jpgs(frames, label, dsize=(224,224), similarity_thresho
frame_filename = "exist-" + frame_filename
cv2.imwrite(frame_filename, frame)
# frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
- frame = cv2.resize(frame, dsize,interpolation=cv2.INTER_LINEAR)
+ frame = cv2.resize(frame, dsize, interpolation=cv2.INTER_LINEAR)
image1_tensor = transform(frame).unsqueeze(0).to(device)
# image1_tensor = (
# torch.tensor(frame, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)
# ).to(device)
- d,i,features=features_model(image1_tensor)
+ d, i, features = features_model(image1_tensor)
if d >= similarity_threshold:
- index_cls[str(i)]=label
+ index_cls[str(i)] = label
else:
index.add(features)
- index_cls[str(index.ntotal-1)]=label
- count+=1
+ index_cls[str(index.ntotal - 1)] = label
+ count += 1
faiss.write_index(index, str(index_path))
- with open(json_path,'w') as f:
- json.dump(index_cls,f)
+ with open(json_path, 'w') as f:
+ json.dump(index_cls, f)
return commitInfo