import logging 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 requests import asyncio import re from qwen_vl_utils import process_vision_info from transformers import AutoModelForVision2Seq, AutoProcessor class qwen_thread: def __init__(self, config,logger): self.config = config self.max_workers = int(config.get("threadnum")) self.executor = ThreadPoolExecutor(max_workers=int(config.get("threadnum"))) self.semaphore = threading.Semaphore(int(config.get("threadnum"))) self.logger = logger # 初始化Milvus集合 connections.connect("default", host=config.get("milvusurl"), port=config.get("milvusport")) # 加载集合 self.collection = Collection(name="smartobject") self.collection.load() if config.get('cuda') is 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("qwenwarning")))] for i in range(int(config.get("qwenwarning"))): model = AutoModelForVision2Seq.from_pretrained( config.get("qwenaddr"), device_map=self.device, trust_remote_code=True, use_safetensors=True, torch_dtype=torch.float16 ).eval() model = model.to(self.device) self.model_pool.append(model) # 共享的处理器 (线程安全) self.processor = AutoProcessor.from_pretrained(config.get("qwenaddr"), use_fast=True) 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_a): try: self.tark_do(res_a, self.config.get("ragurl"), self.config.get("ragmode"), self.config.get("max_tokens")) except Exception as e: self.logger.info(f"处理出错: {e}") raise def tark_do(self,res,ragurl,rag_mode,max_tokens): try: # 生成图片描述 ks_time = datetime.now() risk_description = "" suggestion = "" # 调用规则匹配方法,判断是否预警 is_waning = self.image_rule(res) self.logger.info(f"预警规则规则规则is_waning：{is_waning}") #更新数据的预警结果与数据预警状态 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": 5, #改为已经预警 "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']}") # 保存到milvus image_id = self.collection.insert(data).primary_keys res['id'] = image_id[0] self.logger.info(f"{res['video_point_id']}预警执行完毕：{image_id}运行结束总体用时:{datetime.now() - ks_time}") return None except Exception as e: self.logger.info(f"线程：执行模型解析时出错::{e}") return 0 def image_rule(self, res_data): self.logger.info(f"预警规则规则规则等级分类就是打裂缝多少积分") 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']}?请回答yes或no"}, ], } ] # 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() upper_text = image_des.upper() self.logger.info(f"预警规则规则规则：{upper_text}") if "YES" in upper_text: return 1 else: return 0 except Exception as e: self.logger.info(f"线程：执行图片描述时出错:{e}") return 0 finally: # 4. 释放模型 self._release_model(model) torch.cuda.empty_cache() 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