Merge remote-tracking branch 'origin/main'

This commit is contained in:
Refound-445 2024-11-27 23:09:36 +08:00
commit 2a5872b2c4
10 changed files with 66 additions and 92 deletions

8
.idea/.gitignore generated vendored
View File

@ -1,8 +0,0 @@
# 默认忽略的文件
/shelf/
/workspace.xml
# 基于编辑器的 HTTP 客户端请求
/httpRequests/
# Datasource local storage ignored files
/dataSources/
/dataSources.local.xml

View File

@ -1,6 +0,0 @@
<component name="InspectionProjectProfileManager">
<settings>
<option name="USE_PROJECT_PROFILE" value="false" />
<version value="1.0" />
</settings>
</component>

4
.idea/misc.xml generated
View File

@ -1,4 +0,0 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="ProjectRootManager" version="2" project-jdk-name="Python 3.12" project-jdk-type="Python SDK" />
</project>

8
.idea/modules.xml generated
View File

@ -1,8 +0,0 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="ProjectModuleManager">
<modules>
<module fileurl="file://$PROJECT_DIR$/.idea/nonebot-plugin-nailongremove.iml" filepath="$PROJECT_DIR$/.idea/nonebot-plugin-nailongremove.iml" />
</modules>
</component>
</project>

View File

@ -1,8 +0,0 @@
<?xml version="1.0" encoding="UTF-8"?>
<module type="PYTHON_MODULE" version="4">
<component name="NewModuleRootManager">
<content url="file://$MODULE_DIR$" />
<orderEntry type="inheritedJdk" />
<orderEntry type="sourceFolder" forTests="false" />
</component>
</module>

6
.idea/vcs.xml generated
View File

@ -1,6 +0,0 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="VcsDirectoryMappings">
<mapping directory="" vcs="Git" />
</component>
</project>

View File

@ -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(),

View File

@ -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(

View File

@ -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([
transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5]) # Assuming grayscale or single-channel
])
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()
@ -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)
@ -200,20 +211,23 @@ def process_gif_and_save_jpgs(frames, label, dsize=(224, 224), similarity_thresh
if (
len(
list(
glob.glob(
str(config.nailong_model_dir / "records/*/*.jpg")
),
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"
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