From 2c3655dcfa20d6ac36219b879debef5590f936d4 Mon Sep 17 00:00:00 2001
From: shidong <shidong@jhsoft.cc>
Date: 星期二, 22 七月 2025 14:02:14 +0800
Subject: [PATCH] #2025/7/22 #用vl模型进行预警
---
qwen_thread.py | 127 +++++++++++++++++++++++++++++++++--------
1 files changed, 101 insertions(+), 26 deletions(-)
diff --git a/qwen_thread.py b/qwen_thread.py
index aeb16d9..f1cb6e0 100644
--- a/qwen_thread.py
+++ b/qwen_thread.py
@@ -1,18 +1,15 @@
import time
from concurrent.futures import ThreadPoolExecutor
import threading
-
import torch
from PIL import Image
from pymilvus import connections, Collection
from datetime import datetime
-import os
import requests
import asyncio
-import logging
import re
-from logging.handlers import RotatingFileHandler
+from qwen_vl_utils import process_vision_info
from transformers import AutoModelForVision2Seq, AutoProcessor
@@ -29,19 +26,22 @@
# 鍔犺浇闆嗗悎
self.collection = Collection(name="smartobject")
self.collection.load()
-
+ if config.get('cuda') == None or config.get('cuda') == '0':
+ self.device = f"cuda"
+ else:
+ self.device = f"cuda:{config.get('cuda')}"
self.model_pool = []
self.lock_pool = [threading.Lock() for _ in range(int(config.get("threadnum")))]
for i in range(int(config.get("threadnum"))):
model = AutoModelForVision2Seq.from_pretrained(
config.get("qwenaddr"),
- device_map=f"cuda:{config.get('cuda')}",
+ device_map=self.device,
trust_remote_code=True,
use_safetensors=True,
torch_dtype=torch.float16
).eval()
- model = model.to(f"cuda:{config.get('cuda')}")
+ model = model.to(self.device)
self.model_pool.append(model)
# 鍏变韩鐨勫鐞嗗櫒 (绾跨▼瀹夊叏)
@@ -73,32 +73,58 @@
# 1. 浠庨泦鍚圓鑾峰彇鍚戦噺鍜屽厓鏁版嵁
is_waning = 0
is_desc = 2
-
# 鐢熸垚鍥剧墖鎻忚堪
ks_time = datetime.now()
- desc = self.image_desc(res)
desc_time = datetime.now() - ks_time
current_time = datetime.now()
risk_description = ""
suggestion = ""
+ # 璋冪敤瑙勫垯鍖归厤鏂规硶,鍒ゆ柇鏄惁棰勮
+ is_waning = self.image_rule(res)
+ # 濡傛灉棰勮,鍒欑敓鎴愰殣鎮f弿杩板拰澶勭悊寤鸿
+ if is_waning == 1:
+ # 鑾峰彇瑙勭珷鍒跺害鏁版嵁
+ filedata = self.get_filedata(res['waning_value'],res['suggestion'], ragurl)
+ # 鐢熸垚闅愭偅鎻忚堪
+ risk_description = self.image_rule_chat_with_detail(filedata, res['waning_value'], ragurl,rag_mode,max_tokens)
+ # 鐢熸垚澶勭悊寤鸿
+ suggestion = self.image_rule_chat_suggestion(filedata, res['waning_value'], ragurl,rag_mode,max_tokens)
+ #self.logger.info(f"{res['video_point_id']}鎵ц瀹屾瘯锛歿res['id']}:鏄惁棰勮{is_waning},瀹夊叏闅愭偅锛歿risk_description}\n澶勭悊寤鸿锛歿suggestion}")
+ # 鏁版嵁缁�
+ data = {
+ "event_level_id": res['event_level_id'], # event_level_id
+ "event_level_name": res['event_level_name'], # event_level_id
+ "rule_id": res["rule_id"],
+ "video_point_id": res['video_point_id'], # video_point_id
+ "video_point_name": res['video_point_name'],
+ "is_waning": is_waning,
+ "is_desc": 1,
+ "zh_desc_class": res['zh_desc_class'], # text_vector
+ "bounding_box": res['bounding_box'], # bounding_box
+ "task_id": res['task_id'], # task_id
+ "task_name": res['task_name'], # task_id
+ "detect_id": res['detect_id'], # detect_id
+ "detect_time": res['detect_time'], # detect_time
+ "detect_num": res['detect_num'],
+ "waning_value": res['waning_value'],
+ "image_path": res['image_path'], # image_path
+ "image_desc_path": res['image_desc_path'], # image_desc_path
+ "video_path": res['video_path'],
+ "text_vector": res['text_vector'],
+ "risk_description": risk_description,
+ "suggestion": suggestion,
+ "knowledge_id": res['knowledge_id']
+ }
+ self.collection.delete(f"id == {res['id']}")
+ # 淇濆瓨鍒癿ilvus
+ image_id = self.collection.insert(data).primary_keys
+ res['id'] = image_id[0]
# 鍥剧墖鎻忚堪鐢熸垚鎴愬姛
+ desc = self.image_desc(res)
if desc:
is_desc = 2
- # 璋冪敤瑙勫垯鍖归厤鏂规硶,鍒ゆ柇鏄惁棰勮
- is_waning = self.image_rule_chat(desc, res['waning_value'], ragurl,rag_mode,max_tokens)
- # 濡傛灉棰勮,鍒欑敓鎴愰殣鎮f弿杩板拰澶勭悊寤鸿
- if is_waning == 1:
- # 鑾峰彇瑙勭珷鍒跺害鏁版嵁
- filedata = self.get_filedata(res['waning_value'],res['suggestion'], ragurl)
- # 鐢熸垚闅愭偅鎻忚堪
- risk_description = self.image_rule_chat_with_detail(filedata, res['waning_value'], ragurl,rag_mode,max_tokens)
- # 鐢熸垚澶勭悊寤鸿
- suggestion = self.image_rule_chat_suggestion(filedata, res['waning_value'], ragurl,rag_mode,max_tokens)
- self.logger.info(
- f"{res['video_point_id']}鎵ц瀹屾瘯锛歿res['id']}:鏄惁棰勮{is_waning},瀹夊叏闅愭偅锛歿risk_description}\n澶勭悊寤鸿锛歿suggestion}")
else:
is_desc = 3
-
# 鏁版嵁缁�
data = {
"event_level_id": res['event_level_id'], # event_level_id
@@ -141,7 +167,9 @@
# 璋冪敤rag
asyncio.run(self.insert_json_data(ragurl, data))
rag_time = datetime.now() - current_time
- self.logger.info(f"{res['video_point_id']}鎵ц瀹屾瘯锛歿image_id}杩愯缁撴潫鎬讳綋鐢ㄦ椂:{datetime.now() - ks_time},鍥剧墖鎻忚堪鐢ㄦ椂{desc_time}锛孯AG鐢ㄦ椂{rag_time}")
+ self.logger.info(f"{res['video_point_id']}鎵ц瀹屾瘯锛歿image_id}杩愯缁撴潫鎬讳綋鐢ㄦ椂:{datetime.now() - ks_time},鍥剧墖鎻忚堪鐢ㄦ椂{desc_time},RAG鐢ㄦ椂{rag_time}")
+ if is_waning == 1:
+ self.logger.info(f"{res['video_point_id']}鎵ц瀹屾瘯锛歿image_id},鍥剧墖鎻忚堪锛歿desc}\n闅愭偅锛歿risk_description}\n寤鸿锛歿suggestion}")
except Exception as e:
self.logger.info(f"绾跨▼锛氭墽琛屾ā鍨嬭В鏋愭椂鍑洪敊::{e}")
return 0
@@ -173,12 +201,58 @@
return_tensors="pt",
)
inputs = inputs.to(model.device)
- with torch.inference_mode(),torch.cuda.amp.autocast():
- outputs = model.generate(**inputs,max_new_tokens=200)
+ with torch.inference_mode(), torch.amp.autocast(device_type=self.device, dtype=torch.float16):
+ outputs = model.generate(**inputs,max_new_tokens=200,do_sample=False,num_beams=1,temperature=None,top_p=None,top_k=1,use_cache=True,repetition_penalty=1.0)
generated_ids = outputs[:, len(inputs.input_ids[0]):]
image_text = self.processor.batch_decode(
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
)
+ image_des = (image_text[0]).strip()
+ #self.logger.info(f"{res_data['video_point_id']}:{res_data['id']}:{res_data['detect_time']}:{image_des}")
+ return image_des
+ except Exception as e:
+ self.logger.info(f"绾跨▼锛氭墽琛屽浘鐗囨弿杩版椂鍑洪敊:{e}")
+ finally:
+ # 4. 閲婃斁妯″瀷
+ self._release_model(model)
+ torch.cuda.empty_cache()
+
+ def image_rule(self, res_data):
+ try:
+ model, lock = self._acquire_model()
+ image = Image.open(res_data['image_desc_path']).convert("RGB").resize((600, 600), Image.Resampling.LANCZOS)
+
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image", "image": image},
+ {"type": "text", "text": f"鍥剧墖涓槸鍚︽湁{res_data['waning_value']}?璇峰洖绛攜es鎴杗o"},
+ ],
+ }
+ ]
+
+ # Preparation for inference
+ text = self.processor.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True
+ )
+ image_inputs, video_inputs = process_vision_info(messages)
+ inputs = self.processor(
+ text=[text],
+ images=image_inputs,
+ videos=video_inputs,
+ padding=True,
+ return_tensors="pt",
+ )
+ inputs = inputs.to(model.device)
+
+ with torch.no_grad():
+ outputs = model.generate(**inputs, max_new_tokens=10)
+ generated_ids = outputs[:, len(inputs.input_ids[0]):]
+ image_text = self.processor.batch_decode(
+ generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
+ )
+
image_des = (image_text[0]).strip()
return image_des
except Exception as e:
@@ -239,9 +313,9 @@
}
response = requests.post(ragurl + "/chat", json=search_data)
results = response.json().get('data')
- # self.logger.info(results)
ret = re.sub(r'<think>.*?</think>', '', results, flags=re.DOTALL)
ret = ret.replace(" ", "").replace("\t", "").replace("\n", "")
+ #self.logger.info(f"{rule_text}:{ret}")
is_waning = 0
if len(ret) > 2:
is_waning = 1
@@ -271,6 +345,7 @@
"max_tokens": max_tokens
}
}
+ #self.logger.info(content)
response = requests.post(ragurl + "/chat", json=search_data)
# 浠巎son鎻愬彇data瀛楁鍐呭
ret = response.json()["data"]
--
Gitblit v1.8.0