mirror of
https://github.com/Refound-445/nonebot-plugin-nailongremove.git
synced 2025-11-04 21:22:43 +08:00
up
This commit is contained in:
parent
81c647722a
commit
d98f6bd411
@ -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 黑白名单列表<br />在单级聊天下为该聊天 ID,如 QQ 群号;<br />在多级聊天下为以 `_` 分割的各级聊天 ID,如频道下的子频道或频道下私聊 |
|
||||
| `NAILONG_BLACKLIST` | 否 | `True` | 是否使用黑名单模式 |
|
||||
| `NAILONG_USER_BLACKLIST` | 否 | `[]` | 用户 ID 黑名单列表 |
|
||||
| `NAILONG_PRIORITY` | 否 | `100` | Matcher 优先级 |
|
||||
| **行为配置** | | | |
|
||||
| `NAILONG_RECALL` | 否 | `True` | 是否撤回消息 |
|
||||
| `NAILONG_MUTE_SECONDS` | 否 | `0` | 设置禁言时间,默认为 0 即不禁言<br/>单位:秒 |
|
||||
| `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动图的检测方式<br/>0.检测所有帧<br/>1.只检测第一帧<br/>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 黑白名单列表<br />在单级聊天下为该聊天 ID,如 QQ 群号;<br />在多级聊天下为以 `_` 分割的各级聊天 ID,如频道下的子频道或频道下私聊 |
|
||||
| `NAILONG_BLACKLIST` | 否 | `True` | 是否使用黑名单模式 |
|
||||
| `NAILONG_USER_BLACKLIST` | 否 | `[]` | 用户 ID 黑名单列表 |
|
||||
| `NAILONG_PRIORITY` | 否 | `100` | Matcher 优先级 |
|
||||
| **行为配置** | | | |
|
||||
| `NAILONG_RECALL` | 否 | `True` | 是否撤回消息 |
|
||||
| `NAILONG_MUTE_SECONDS` | 否 | `0` | 设置禁言时间,默认为 0 即不禁言<br/>单位:秒 |
|
||||
| `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),可用变量见下,可以根据标签自定义对应值,随机发送列表其中一条消息,<br/>如遇其中没有的标签会回退到 `nailong`<br/>如果对应值为空列表`[]`,则会检测而不会发送消息 |
|
||||
| `NAILONG_FAILED_TIP` | 否 | `{"nailong": ["{:Reply($message_id)}呜,不要发奶龙了嘛 🥺 👉👈"]}` | 撤回失败或禁用撤回时发送的提示,同上 |
|
||||
| `NAILONG_CHECK_ALL_FRAMES` | 否 | `False` | 使用模型 1 时是否检查图片中的所有帧,需要同时设置`NAILONG_CHECK_MODE`为0,启用该项后消息模板中的 `$checked_result` 变量当原图为动图时会变为动图 |
|
||||
| `NAILONG_CHECK_MODE` | 否 | `0` | 选择对GIF动图的检测方式<br/>0.检测所有帧<br/>1.只检测第一帧<br/>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`自行选择
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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(),
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user