diff --git a/packages/nonebot-plugin-nailongremove-base/README.md b/packages/nonebot-plugin-nailongremove-base/README.md
index 51d18da..c8e00f2 100644
--- a/packages/nonebot-plugin-nailongremove-base/README.md
+++ b/packages/nonebot-plugin-nailongremove-base/README.md
@@ -57,7 +57,7 @@ NaiLongRemove 是一款由简单的 AI 模型建立的奶龙识别插件,可
### 技术
-目前插件支持两种模型,可通过配置文件更换,详见文档下方配置一节。
+目前插件支持三种模型,可通过配置文件更换,详见文档下方配置一节。
用户可以根据需要自行选择心仪的模型,两个模型性能都已经经过优化,但仍可能会有不同程度的误差,也欢迎各位继续反馈给我们~
## 💿 安装
@@ -187,44 +187,44 @@ pip install nonebot-plugin-nailongremove-base -U
在 nonebot2 项目的 `.env` 文件中添加下表中的必填配置
-| 配置项 | 必填 | 默认值 | 说明 |
-|:-----------------------------------:| :--: |:-------------------------------------------------------:|:----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------:|
-| **全局配置** | | | |
-| `PROXY` | 否 | `None` | 下载模型等文件时使用的代理地址 |
-| **响应配置** | | | |
-| `NAILONG_BYPASS_SUPERUSER` | 否 | `True` | 是否不检查超级用户发送的图片 |
-| `NAILONG_BYPASS_ADMIN` | 否 | `True` | 是否不检查群组管理员发送的图片 |
-| `NAILONG_NEED_ADMIN` | 否 | `False` | 当自身不为群组管理员时是否不检查群内所有图片 |
-| `NAILONG_LIST_SCENES` | 否 | `[]` | 聊天场景 ID 黑白名单列表
在单级聊天下为该聊天 ID,如 QQ 群号;
在多级聊天下为以 `_` 分割的各级聊天 ID,如频道下的子频道或频道下私聊 |
-| `NAILONG_BLACKLIST` | 否 | `True` | 是否使用黑名单模式 |
-| `NAILONG_USER_BLACKLIST` | 否 | `[]` | 用户 ID 黑名单列表 |
-| `NAILONG_PRIORITY` | 否 | `100` | Matcher 优先级 |
-| **行为配置** | | | |
-| `NAILONG_RECALL` | 否 | `True` | 是否撤回消息 |
-| `NAILONG_MUTE_SECONDS` | 否 | `0` | 设置禁言时间,默认为 0 即不禁言
单位:秒 |
-| `NAILONG_TIP` | 否 | `{"nailong": "本群禁止发奶龙!"}` | 发送的提示,使用 [Alconna 的消息模板](https://nonebot.dev/docs/best-practice/alconna/uniseg#%E4%BD%BF%E7%94%A8%E6%B6%88%E6%81%AF%E6%A8%A1%E6%9D%BF),可用变量见下,可以根据标签自定义对应值,如遇其中没有的标签会回退到 `nailong` |
-| `NAILONG_FAILED_TIP` | 否 | `{"nailong": "{:Reply($message_id)}呜,不要发奶龙了嘛 🥺 👉👈"}` | 撤回失败或禁用撤回时发送的提示,同上 |
-| `NAILONG_CHECK_ALL_FRAMES` | 否 | `False` | 使用模型 1 时是否检查图片中的所有帧,需要同时设置`NAILONG_CHECK_MODE`为0,启用该项后消息模板中的 `$checked_result` 变量当原图为动图时会变为动图 |
-| `NAILONG_CHECK_MODE` | 否 | `0` | 选择对GIF动图的检测方式
0.检测所有帧
1.只检测第一帧
2.随机抽帧检测 |
-| **相似度检测配置** | | | |
-| `NAILONG_SIMILARITY_ON` | 否 | `False` | 是否启用处理图片前对本地存储进行相似度检测(该功能仍在更新中,目前可能耗能较大且处理较慢) |
-| `NAILONG_SIMILARITY_MAX_STORAGE` | 否 | `10` | 本地存储报错图片最大上限,到达上限会压缩并删除上次记录 |
-| `NAILONG_SIMILARITY_MAX_BATCH_SIZE` | 否 | `10` | 本地存储相似度检测时处理的最大批数量 |
-| **模型通用配置** | | | |
-| `NAILONG_MODEL_DIR` | 否 | `./data/nailongremove` | 模型的下载位置 |
-| `NAILONG_MODEL` | 否 | `1` | 选择需要加载的模型,可用模型见下 |
-| `NAILONG_AUTO_UPDATE_MODEL` | 否 | `True` | 是否自动更新模型 |
-| `NAILONG_CONCURRENCY` | 否 | `1` | 当图片为动图时,针对该图片并发识别图片帧的最大并发数 |
-| `NAILONG_ONNX_PROVIDERS` | 否 | `["CPUExecutionProvider"]` | 加载 onnx 模型使用的 provider 列表,请参考上方安装文档 |
-| **模型 1 特定配置** | | | |
-| `NAILONG_MODEL1_TYPE` | 否 | `tiny` | 模型 1 使用的模型类型,可用 `tiny` / `m` |
-| `NAILONG_MODEL1_YOLOX_SIZE` | 否 | `None` | 针对模型 1,自定义模型输入可能会有尺寸更改 |
-| **模型 2 特定配置** | | | |
-| `NAILONG_MODEL2_ONLINE` | 否 | `False` | 针对模型 2,是否启用在线推理,此模式目前不适用`NAILONG_CHECK_MODE`为0 |
-| **模型 1&2 特定配置** | | | |
-| `NAILONG_MODEL1_SCORE` | 否 | `{"nailong": 0.5}` | 模型 1&2 置信度阈值,范围 `0` ~ `1`,可以根据标签自定义对应值,设置对应标签的阈值以检测该标签,设为 `null` 或者不填可以忽略该标签 |
-| **杂项配置** | | | |
-| `NAILONG_GITHUB_TOKEN` | 否 | `None` | GitHub Access Token,遇到模型下载或更新问题时可尝试填写 |
+| 配置项 | 必填 | 默认值 | 说明 |
+|:--------------------------------:|:--:|:---------------------------------------------------------:|:----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------:|
+| **全局配置** | | | |
+| `PROXY` | 否 | `None` | 下载模型等文件时使用的代理地址 |
+| **响应配置** | | | |
+| `NAILONG_BYPASS_SUPERUSER` | 否 | `True` | 是否不检查超级用户发送的图片 |
+| `NAILONG_BYPASS_ADMIN` | 否 | `True` | 是否不检查群组管理员发送的图片 |
+| `NAILONG_NEED_ADMIN` | 否 | `False` | 当自身不为群组管理员时是否不检查群内所有图片 |
+| `NAILONG_LIST_SCENES` | 否 | `[]` | 聊天场景 ID 黑白名单列表
在单级聊天下为该聊天 ID,如 QQ 群号;
在多级聊天下为以 `_` 分割的各级聊天 ID,如频道下的子频道或频道下私聊 |
+| `NAILONG_BLACKLIST` | 否 | `True` | 是否使用黑名单模式 |
+| `NAILONG_USER_BLACKLIST` | 否 | `[]` | 用户 ID 黑名单列表 |
+| `NAILONG_PRIORITY` | 否 | `100` | Matcher 优先级 |
+| **行为配置** | | | |
+| `NAILONG_RECALL` | 否 | `True` | 是否撤回消息 |
+| `NAILONG_MUTE_SECONDS` | 否 | `0` | 设置禁言时间,默认为 0 即不禁言
单位:秒 |
+| `NAILONG_TIP` | 否 | `{"nailong": ["本群禁止发奶龙!"]}` | 发送的提示,使用 [Alconna 的消息模板](https://nonebot.dev/docs/best-practice/alconna/uniseg#%E4%BD%BF%E7%94%A8%E6%B6%88%E6%81%AF%E6%A8%A1%E6%9D%BF),可用变量见下,可以根据标签自定义对应值,随机发送列表其中一条消息,
如遇其中没有的标签会回退到 `nailong`
如果对应值为空列表`[]`,则会检测而不会发送消息 |
+| `NAILONG_FAILED_TIP` | 否 | `{"nailong": ["{:Reply($message_id)}呜,不要发奶龙了嘛 🥺 👉👈"]}` | 撤回失败或禁用撤回时发送的提示,同上 |
+| `NAILONG_CHECK_ALL_FRAMES` | 否 | `False` | 使用模型 1 时是否检查图片中的所有帧,需要同时设置`NAILONG_CHECK_MODE`为0,启用该项后消息模板中的 `$checked_result` 变量当原图为动图时会变为动图 |
+| `NAILONG_CHECK_MODE` | 否 | `0` | 选择对GIF动图的检测方式
0.检测所有帧
1.只检测第一帧
2.随机抽帧检测 |
+| **相似度检测配置** | | | |
+| `NAILONG_SIMILARITY_ON` | 否 | `False` | 是否启用处理图片前对本地存储进行相似度检测(该功能仍在更新中,目前可能耗能较大且处理较慢) |
+| `NAILONG_SIMILARITY_MAX_STORAGE` | 否 | `10` | 本地存储报错图片最大上限,到达上限会压缩并删除上次记录 |
+| `NAILONG_HF_TOKEN` | 否 | `None` | Hugging Face Access Token,自动上传数据到hf,并成为数据集贡献者 |
+| **模型通用配置** | | | |
+| `NAILONG_MODEL_DIR` | 否 | `./data/nailongremove` | 模型的下载位置 |
+| `NAILONG_MODEL` | 否 | `1` | 选择需要加载的模型,可用模型见下 |
+| `NAILONG_AUTO_UPDATE_MODEL` | 否 | `True` | 是否自动更新模型 |
+| `NAILONG_CONCURRENCY` | 否 | `1` | 当图片为动图时,针对该图片并发识别图片帧的最大并发数 |
+| `NAILONG_ONNX_PROVIDERS` | 否 | `["CPUExecutionProvider"]` | 加载 onnx 模型使用的 provider 列表,请参考上方安装文档 |
+| **模型 1 特定配置** | | | |
+| `NAILONG_MODEL1_TYPE` | 否 | `tiny` | 模型 1 使用的模型类型,可用 `tiny` / `m` |
+| `NAILONG_MODEL1_YOLOX_SIZE` | 否 | `None` | 针对模型 1,自定义模型输入可能会有尺寸更改 |
+| **模型 2 特定配置** | | | |
+| `NAILONG_MODEL2_ONLINE` | 否 | `False` | 针对模型 2,是否启用在线推理,此模式目前不适用`NAILONG_CHECK_MODE`为0 |
+| **模型 1&2 特定配置** | | | |
+| `NAILONG_MODEL1_SCORE` | 否 | `{"nailong": 0.5}` | 模型 1&2 置信度阈值,范围 `0` ~ `1`,可以根据标签自定义对应值,设置对应标签的阈值以检测该标签,设为 `null` 或者不填可以忽略该标签 |
+| **杂项配置** | | | |
+| `NAILONG_GITHUB_TOKEN` | 否 | `None` | GitHub Access Token,遇到模型下载或更新问题时可尝试填写 |
### 可用模型
@@ -259,6 +259,12 @@ pip install nonebot-plugin-nailongremove-base -U
## 📝 更新日志
+### 2.3.3
+
+- 优化临时处理方案,减小性能压力同时提升速度(向量库faiss也支持GPU处理,但非专业人士不推荐使用GPU,因为这个安装过程比较复杂)
+- 增加`NAILONG_HF_TOKEN`实现自动将报错图片上传Hugging Face数据集
+- 更改配置项`NAILONG_TIP`和`NAILONG_FAILED_TIP`格式,允许随机发送返回消息,并且对应值为空列表`[]`时,仅检测图片(或者禁言撤回)而不会返回消息
+
### 2.3.2
- 更新对GIF动图的三种帧处理模式,通过`NAILONG_CHECK_MODE`自行选择
diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/config.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/config.py
index 3e272b0..d7a29b9 100644
--- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/config.py
+++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/config.py
@@ -41,11 +41,11 @@ class Config(BaseModel):
nailong_recall: bool = True
nailong_mute_seconds: int = 0
- nailong_tip: Dict[str, str] = {
- DEFAULT_LABEL: "本群禁止发送奶龙!",
+ nailong_tip: Dict[str, List[str]] = {
+ DEFAULT_LABEL: ["本群禁止发送奶龙!"],
}
- nailong_failed_tip: Dict[str, str] = {
- DEFAULT_LABEL: "{:Reply($message_id)}呜,不要发奶龙了嘛 🥺 👉👈",
+ nailong_failed_tip: Dict[str, List[str]] = {
+ DEFAULT_LABEL: ["{:Reply($message_id)}呜,不要发奶龙了嘛 🥺 👉👈"],
}
nailong_check_all_frames: bool = False
@@ -65,8 +65,8 @@ class Config(BaseModel):
nailong_model2_online: bool = False
nailong_check_mode: int = 0
nailong_similarity_on: bool = False
- nailong_similarity_max_storage: int = 10
- nailong_similarity_max_batch_size: int = 10
+ nailong_similarity_max_storage: int = 1000
+ nailong_hf_token: Optional[str] = None
nailong_github_token: Optional[str] = None
diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/handler.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/handler.py
index fe24f64..178a792 100644
--- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/handler.py
+++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/handler.py
@@ -1,3 +1,4 @@
+import random
import re
from typing import Any, Awaitable, Callable, Iterable, List, TypeVar
@@ -104,14 +105,14 @@ async def handle_function(bot: BaseBot, ev: BaseEvent, msg: UniMsg, session: Uni
frames.append(temp_image)
except StopIteration:
break
- zip_filename = process_gif_and_save_jpgs(frames, label, input_shape)
- if zip_filename is None:
+ commitInfo = process_gif_and_save_jpgs(frames, label, (224,224))
+ if commitInfo is None:
await nailong.finish(
- f"已保存数据到目录{config.nailong_model_dir}\\records\\{label},标签:{label}",
+ f"The new data has been saved to the directory {config.nailong_model_dir}\\records\\{label}, label: {label}.",
)
else:
await nailong.finish(
- f"记录数据超过{config.nailong_similarity_max_storage},已清除原记录数据,压缩并保存至{zip_filename}\n已保存数据到目录{config.nailong_model_dir}\\records\\{label},标签:{label}",
+ f"The recorded data has exceeded {config.nailong_similarity_max_storage}, the original data has been cleared, compressed, and upload to {commitInfo.commit_url}\nThe new data has been saved to the directory {config.nailong_model_dir}\\records\\{label}, label: {label}.",
)
else:
try:
@@ -131,9 +132,12 @@ async def handle_function(bot: BaseBot, ev: BaseEvent, msg: UniMsg, session: Uni
template_dict = (
config.nailong_tip if punish_ok else config.nailong_failed_tip
)
- template_str = template_dict[
+ template_str_all = template_dict[
check_res.label if (check_res.label in template_dict) else DEFAULT_LABEL
]
+ if len(template_str_all) == 0:
+ continue
+ template_str=template_str_all[random.randint(0, len(template_str_all) - 1)]
mapping = {
"$event": ev,
"$target": msg.get_target(),
diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/classification.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/classification.py
index a45bf0a..c4c4b64 100644
--- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/classification.py
+++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/classification.py
@@ -37,7 +37,7 @@ SIZE = 224
@run_sync
def check_single(image: np.ndarray, is_gif: bool = False) -> CheckSingleResult[None]:
if is_gif:
- res = similarity_process(image, dsize=(SIZE, SIZE))
+ res = similarity_process(image)
if res is not None:
return res
return CheckSingleResult.not_ok(None)
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 970ba21..0e95ae0 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
@@ -68,7 +68,7 @@ input_shape = config.nailong_model1_yolox_size or config.nailong_model1_type.yol
@run_sync
def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult:
if is_gif:
- res = similarity_process(frame, dsize=input_shape)
+ res = similarity_process(frame)
if res is not None:
return CheckSingleResult(ok=res.ok, label=res.label, extra=frame)
return CheckSingleResult(ok=False, label=None, extra=frame)
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 d0a559e..ae4f92f 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
@@ -87,7 +87,7 @@ def _check_single(
is_gif: bool = False,
) -> CheckSingleResult[Optional[Detections]]:
if is_gif:
- res = similarity_process(frame, dsize=input_shape)
+ res = similarity_process(frame)
if res is not None:
return res
return CheckSingleResult.not_ok(None)
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 6d7bc4e..ea4a3c8 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
@@ -7,13 +7,13 @@ 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
@@ -21,6 +21,60 @@ 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
+ from torchvision import transforms
+
+ transform = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize(mean=[0.5], std=[0.5]) # Assuming grayscale or single-channel
+ ])
+ class MyModel(
+ nn.Module,
+ PyTorchModelHubMixin,
+ ):
+ 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
+
+ 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'
+ 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:
+ index_cls=json.load(f)
+ else:
+ 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]
+ 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
+ features_model.resnet.fc.register_forward_hook(hook)
+ features_model.eval()
+
+
@dataclass
class CheckSingleResult(Generic[T]):
@@ -121,113 +175,79 @@ async def race_check(
return None
-def similarity_process(image1: np.ndarray, dsize) -> Optional[CheckSingleResult]:
- path = list(glob.glob(os.path.join(config.nailong_model_dir, "records/*/*.jpg")))
- if len(path) == 0:
- return None
- image1 = cv2.cvtColor(image1, cv2.COLOR_BGR2RGB)
+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 = (
- torch.tensor(image1, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)
- )
- image1_tensor = image1_tensor.reshape(1, -1).to(device)
- for i in range(0, len(path), config.nailong_similarity_max_batch_size):
- temp_paths = path[
- i : (min(len(path), i + config.nailong_similarity_max_batch_size))
- ]
- image2s = []
- for image_path in temp_paths:
- image2 = cv2.imread(image_path)
- image2 = cv2.cvtColor(image2, cv2.COLOR_BGR2RGB)
- image2 = cv2.resize(image2, dsize, interpolation=cv2.INTER_LINEAR)
- image2s.append(image2)
- image2_tensor = torch.tensor(np.array(image2s), dtype=torch.float32).permute(
- 0,
- 3,
- 1,
- 2,
- )
- image2_tensor = image2_tensor.reshape(image2_tensor.shape[0], -1).to(device)
- similarities = F.cosine_similarity(image1_tensor, image2_tensor)
- indices = torch.nonzero(similarities > 0.99)
- index = indices[0].item() if indices.numel() > 0 else None
- if index is not None:
- image_path = path[index]
- label = os.path.split(image_path)[-2].split("\\")[-1]
- return CheckSingleResult(ok=True, label=label, extra=None)
+ 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)
+ if distance >= similarity_threshold:
+ label =index_cls[str(indice)]
+ return CheckSingleResult(ok=True, label=label, extra=None)
return None
-def process_gif_and_save_jpgs(frames, label, dsize, similarity_threshold=0.85):
+def process_gif_and_save_jpgs(frames, label, dsize=(224,224), similarity_threshold=1):
if (
len(
list(
glob.glob(
- os.path.join(str(config.nailong_model_dir), "records/*/*.jpg"),
+ str(config.nailong_model_dir / "records/*/*.jpg")
),
),
)
- >= config.nailong_similarity_max_storage
+ >= config.nailong_similarity_max_storage and config.nailong_hf_token is not None
):
zip_filename = shutil.make_archive(
- os.path.join(
- str(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",
- os.path.join(str(config.nailong_model_dir), "records"),
+ config.nailong_model_dir / "records"
)
- shutil.rmtree(os.path.join(str(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,
+ path_in_repo="new_dataset.zip",
+ repo_id="refoundd/NailongClassification",
+ repo_type="dataset",
+ create_pr=True,
+ token=config.nailong_hf_token,
+ )
+ # os.remove(zip_filename)
else:
- zip_filename = None
- output_dir = os.path.join(str(config.nailong_model_dir), "records", label)
+ commitInfo = None
+ output_dir = config.nailong_model_dir / "records"/ label
if not os.path.exists(output_dir):
os.makedirs(output_dir)
- frame_count = [i for i in range(len(frames))]
- while len(frame_count) > 0:
- frame_num1 = frame_count[0]
- frame_count.remove(frame_num1)
- frame1 = frames[frame_num1]
+ count=0
+ for frame in frames:
frame_filename = os.path.join(
output_dir,
"frame{}_{}.jpg".format(
- frame_num1,
+ count,
datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
),
)
while os.path.exists(frame_filename):
frame_filename = "exist-" + frame_filename
- cv2.imwrite(frame_filename, frame1)
- # frame1 = cv2.cvtColor(frame1, cv2.COLOR_BGR2RGB)
- frame1 = cv2.resize(frame1, dsize)
- image1_tensor = (
- torch.tensor(frame1, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)
- )
- image1_tensor = image1_tensor.reshape(1, -1).to(device)
- max_length = len(list(frame_count))
- indexs = []
- for i in range(0, max_length, config.nailong_similarity_max_batch_size):
- frame2_num = frame_count[
- i : (min(max_length, i + config.nailong_similarity_max_batch_size))
- ]
- frame2 = [frames[i] for i in frame2_num]
- # frame2 = cv2.cvtColor(frame2, cv2.COLOR_BGR2RGB)
- frame2 = [cv2.resize(t, dsize) for t in frame2]
- image2_tensor = torch.tensor(np.array(frame2), dtype=torch.float32).permute(
- 0,
- 3,
- 1,
- 2,
- )
- image2_tensor = image2_tensor.reshape(image2_tensor.shape[0], -1).to(device)
- similarities = F.cosine_similarity(image1_tensor, image2_tensor)
- indices = torch.nonzero(similarities > similarity_threshold)
- index = indices.squeeze().tolist() if indices.numel() > 0 else None
- if type(index) is int:
- index = [index]
- if index is not None:
- indexs.extend([frame2_num[i] for i in index])
- frame_count = [i for i in frame_count if i not in indexs]
- return zip_filename
+ cv2.imwrite(frame_filename, frame)
+ # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+ 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)
+ if d >= similarity_threshold:
+ index_cls[str(i)]=label
+ else:
+ index.add(features)
+ 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)
+ return commitInfo
diff --git a/packages/nonebot-plugin-nailongremove-base/pyproject.toml b/packages/nonebot-plugin-nailongremove-base/pyproject.toml
index fa3d399..8814b1b 100644
--- a/packages/nonebot-plugin-nailongremove-base/pyproject.toml
+++ b/packages/nonebot-plugin-nailongremove-base/pyproject.toml
@@ -22,6 +22,9 @@ dependencies = [
"huggingface-hub>=0.26.2",
"ultralytics>=8.3.31",
"gradio-client>=1.3.0",
+ "faiss-cpu>=1.9.0.post1",
+ "faiss-gpu>=1.7.2",
+ "scikit-learn>=1.5.2",
]
license = { text = "MIT" }
readme = "README.md"