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 transformers import AutoModelForVision2Seq, AutoProcessor class qwen_thread: def __init__(self, max_workers,config,model_path): self.executor = ThreadPoolExecutor(max_workers=max_workers) self.semaphore = threading.Semaphore(max_workers) self.max_workers = max_workers # 初始化Milvus集合 connections.connect("default", host=config.get("milvusurl"), port=config.get("milvusport")) # 加载集合 self.collection = Collection(name="smartobject") self.collection.load() self.config = config self.model_pool = [] self.lock_pool = [threading.Lock() for _ in range(max_workers)] for i in range(max_workers): model = AutoModelForVision2Seq.from_pretrained( model_path, device_map="cuda:1", trust_remote_code=True, torch_dtype=torch.float16 ).eval() self.model_pool.append(model) # 共享的处理器 (线程安全) self.processor = AutoProcessor.from_pretrained(model_path,use_fast=True) # 创建实例专属logger self.logger = logging.getLogger(f"{self.__class__}_{id(self)}") self.logger.setLevel(logging.INFO) # 避免重复添加handler if not self.logger.handlers: handler = RotatingFileHandler( filename=os.path.join("logs", 'thread_log.log'), maxBytes=10 * 1024 * 1024, backupCount=3, encoding='utf-8' ) formatter = logging.Formatter( '%(asctime)s - %(filename)s:%(lineno)d - %(funcName)s() - %(levelname)s: %(message)s' ) handler.setFormatter(formatter) self.logger.addHandler(handler) def submit(self,res_a): # 尝试获取信号量（非阻塞） acquired = self.semaphore.acquire(blocking=False) if not acquired: self.logger.info(f"线程池已满，等待空闲线程... (当前活跃: {self.max_workers - self.semaphore._value}/{self.max_workers})") # 阻塞等待直到有可用线程 self.semaphore.acquire(blocking=True) future = self.executor.submit(self._wrap_task, res_a) future.add_done_callback(self._release_semaphore) return future def _wrap_task(self, res): try: #self.logger.info(f"处理: { res['id']}开始") current_time = datetime.now() image_id = self.tark_do(res, self.config.get("ragurl"), self.config.get("ragmode"), self.config.get("max_tokens")) self.logger.info(f"处理: { res['id']}完毕{image_id}:{datetime.now() - current_time}") return image_id except Exception as e: self.logger.info(f"任务 { res['id']} 处理出错: {e}") raise def tark_do(self,res,ragurl,rag_mode,max_tokens): try : # 1. 从集合A获取向量和元数据 is_waning = 0 image_des = self.image_desc(f"{res['image_desc_path']}") self.logger.info(image_des) if image_des: # rule_text = self.get_rule(ragurl) is_waning = self.image_rule_chat(image_des,res['waning_value'],ragurl,rag_mode,max_tokens) is_desc = 2 else: is_waning = 0 is_desc = 3 data = { "id": res['id'], "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": is_desc, "zh_desc_class": image_des, # 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'] } # 保存到milvus image_id = self.collection.upsert(data).primary_keys if is_desc == 2: data = { "id": str(image_id[0]), "video_point_id": res['video_point_id'], "video_path": res["video_point_name"], "zh_desc_class": image_des, "detect_time": res['detect_time'], "image_path": f"{res['image_path']}", "task_name": res["task_name"], "event_level_name": res["event_level_name"], "rtsp_address": f"{res['video_path']}" } # 调用rag asyncio.run(self.insert_json_data(ragurl, data)) return image_id except Exception as e: self.logger.info(f"线程：执行模型解析时出错:任务：{e}") return 0 def image_desc(self, image_path): try: model, lock = self._acquire_model() # 2. 处理图像 image = Image.open(image_path).convert("RGB") # 替换为您的图片路径 image = image.resize((600, 600), Image.Resampling.LANCZOS) # 高质量缩放 messages = [ { "role": "user", "content": [ { "type": "image", }, {"type": "text", "text": "请详细描述图片中的目标信息及特征。返回格式为整段文字描述"}, ], } ] # Preparation for inference text = self.processor.apply_chat_template( messages, add_generation_prompt=True ) inputs = self.processor( text=[text], images=[image], padding=True, return_tensors="pt", ) inputs = inputs.to("cuda:1") current_time = datetime.now() outputs = model.generate(**inputs, max_new_tokens=300, do_sample=True, temperature=0.7, renormalize_logits=True ) print(f"处理完毕:{datetime.now() - current_time}") 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: self.logger.info(f"线程：执行图片描述时出错:{e}") finally: # 4. 释放模型 self._release_model(model) torch.cuda.empty_cache() def get_rule(self,ragurl): try: rule_text = None search_data = { "collection_name": "smart_rule", "query_text": "", "search_mode": "hybrid", "limit": 100, "weight_dense": 0.7, "weight_sparse": 0.3, "filter_expr": "", "output_fields": ["text"] } response = requests.post(ragurl + "/search", json=search_data) results = response.json().get('results') rule_text = "" ruleid = 1 for rule in results: if rule['score'] >= 0: rule_text = rule_text + str(ruleid) + ". " + rule['entity'].get('text') + ";\n" ruleid = ruleid + 1 # self.logger.info(len(rule_text)) else: self.logger.info(f"线程：执行获取规则时出错:{response}") return rule_text except Exception as e: self.logger.info(f"线程：执行获取规则时出错:{e}") return None def image_rule_chat(self, image_des,rule_text, ragurl, rag_mode,max_tokens): try: content = ( f"图片描述内容为：\n{image_des}\n规则内容：\n{rule_text}。\n请验证图片描述中是否有符合规则的内容，不进行推理和think。返回结果格式为[xxx符合的规则id]，如果没有返回[]") # self.logger.info(content) #self.logger.info(len(content)) search_data = { "prompt": "", "messages": [ { "role": "user", "content": content } ], "llm_name": rag_mode, "stream": False, "gen_conf": { "temperature": 0.7, "max_tokens": max_tokens } } response = requests.post(ragurl + "/chat", json=search_data) results = response.json().get('data') #self.logger.info(len(results)) # self.logger.info(results) ret = re.sub(r'.*?', '', results, flags=re.DOTALL) ret = ret.replace(" ", "").replace("\t", "").replace("\n", "") is_waning = 0 if len(ret) > 2: is_waning = 1 return is_waning except Exception as e: self.logger.info(f"线程：执行规则匹配时出错:{image_des, rule_text, ragurl, rag_mode,e}") return None async def insert_json_data(self, ragurl, data): try: data = {'collection_name': "smartrag", "data": data, "description": ""} requests.post(ragurl + "/insert_json_data", json=data, timeout=(0.3, 0.3)) #self.logger.info(f"调用录像服务:{ragurl, data}") except Exception as e: #self.logger.info(f"{self._thread_name}线程：调用录像时出错:地址：{ragurl}：{e}") return def _release_semaphore(self, future): self.semaphore.release() self.logger.info(f"释放线程 (剩余空闲: {self.semaphore._value}/{self.max_workers})") def shutdown(self): """安全关闭""" self.executor.shutdown(wait=False) for model in self.model_pool: del model torch.cuda.empty_cache() def _acquire_model(self): """从池中获取一个空闲模型 (简单轮询)""" while True: for i, (model, lock) in enumerate(zip(self.model_pool, self.lock_pool)): if lock.acquire(blocking=False): return model, lock time.sleep(0.1) # 避免CPU空转 def _release_model(self, model): """释放模型回池""" for i, m in enumerate(self.model_pool): if m == model: self.lock_pool[i].release() break def remove_duplicate_lines(self,text): seen = set() result = [] for line in text.split('。'): # 按句号分割 if line.strip() and line not in seen: seen.add(line) result.append(line) return '。'.join(result) def remove_duplicate_lines_d(self,text): seen = set() result = [] for line in text.split(','): # 按句号分割 if line.strip() and line not in seen: seen.add(line) result.append(line) return '。'.join(result) def remove_duplicate_lines_n(self,text): seen = set() result = [] for line in text.split('\n'): # 按句号分割 if line.strip() and line not in seen: seen.add(line) result.append(line) return '。'.join(result)