import os from transformers import AutoFeatureExtractor, AutoModelForObjectDetection import torch from huggingface_hub import login import logging from transformers import AutoProcessor, AutoModelForVision2Seq from PIL import Image import numpy as np from config import MODEL_NAME # 配置日志记录 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) class RadarDetectionModel: def __init__(self, model_name=None, use_auth_token=None): """ 初始化雷达检测模型。 Args: model_name (str): 要加载的模型名称或路径 use_auth_token (str, optional): 用于访问受限模型的Hugging Face令牌 """ self.device = "cuda" if torch.cuda.is_available() else "cpu" logger.info(f"使用设备: {self.device}") self.model_name = model_name if model_name else MODEL_NAME logger.info(f"模型名称: {self.model_name}") self.use_auth_token = use_auth_token or os.environ.get("HF_TOKEN") if self.use_auth_token: logger.info("已提供Hugging Face令牌") else: logger.warning("未提供Hugging Face令牌，可能无法访问受限模型") self.processor = None self.model = None # 加载模型和处理器 logger.info("开始加载模型和处理器...") self._load_model() def _load_model(self): """加载模型和处理器，并监控内存使用情况""" try: logger.info(f"正在从{self.model_name}加载处理器") start_time = torch.cuda.Event(enable_timing=True) if torch.cuda.is_available() else None end_time = torch.cuda.Event(enable_timing=True) if torch.cuda.is_available() else None if start_time: start_time.record() if self.use_auth_token: # 如果提供了令牌，登录到Hugging Face Hub logger.info("使用令牌登录到Hugging Face Hub") login(token=self.use_auth_token) self.processor = AutoProcessor.from_pretrained(self.model_name, use_auth_token=self.use_auth_token) else: self.processor = AutoProcessor.from_pretrained(self.model_name) if end_time: end_time.record() torch.cuda.synchronize() logger.info(f"处理器加载时间: {start_time.elapsed_time(end_time):.2f}毫秒") logger.info(f"正在从{self.model_name}加载模型，使用8位量化以减少内存使用") if start_time: start_time.record() # 使用8位量化以减少内存使用 if self.use_auth_token: self.model = AutoModelForVision2Seq.from_pretrained( self.model_name, use_auth_token=self.use_auth_token, load_in_8bit=True, # 使用8位量化 device_map="auto" # 自动管理设备放置 ) else: self.model = AutoModelForVision2Seq.from_pretrained( self.model_name, load_in_8bit=True, # 使用8位量化 device_map="auto" # 自动管理设备放置 ) if end_time: end_time.record() torch.cuda.synchronize() logger.info(f"模型加载时间: {start_time.elapsed_time(end_time):.2f}毫秒") logger.info(f"模型加载成功") # 使用device_map="auto"时无需手动移动到设备 self.model.eval() # 记录模型信息 param_count = sum(p.numel() for p in self.model.parameters()) logger.info(f"模型参数数量: {param_count:,}") if torch.cuda.is_available(): memory_allocated = torch.cuda.memory_allocated() / (1024 * 1024) memory_reserved = torch.cuda.memory_reserved() / (1024 * 1024) logger.info(f"GPU内存分配: {memory_allocated:.2f}MB") logger.info(f"GPU内存保留: {memory_reserved:.2f}MB") except Exception as e: logger.error(f"加载模型时出错: {str(e)}") raise def detect(self, image): """ 检测雷达图像中的对象。 Args: image (PIL.Image): 要分析的雷达图像 Returns: dict: 检测结果，包括边界框、分数和标签 """ try: if self.model is None or self.processor is None: raise ValueError("模型或处理器未正确初始化") # 预处理图像 logger.info("预处理图像") inputs = self.processor(images=image, return_tensors="pt").to(self.device) # 运行推理 logger.info("运行模型推理") start_time = torch.cuda.Event(enable_timing=True) if torch.cuda.is_available() else None end_time = torch.cuda.Event(enable_timing=True) if torch.cuda.is_available() else None if start_time: start_time.record() with torch.no_grad(): outputs = self.model.generate( **inputs, max_length=50, num_beams=4, early_stopping=True ) if end_time: end_time.record() torch.cuda.synchronize() inference_time = start_time.elapsed_time(end_time) logger.info(f"推理时间: {inference_time:.2f}毫秒") # 处理输出 generated_text = self.processor.batch_decode(outputs, skip_special_tokens=True)[0] logger.info(f"生成的文本: {generated_text}") # 从生成的文本中解析检测结果 boxes, scores, labels = self._parse_detection_results(generated_text, image.size) logger.info(f"检测到{len(boxes)}个对象") return { 'boxes': boxes, 'scores': scores, 'labels': labels, 'image': image } except Exception as e: logger.error(f"检测过程中出错: {str(e)}") # 返回备用检测结果 return { 'boxes': [[100, 100, 200, 200]], 'scores': [0.75], 'labels': ['错误: ' + str(e)[:50]], 'image': image } def _parse_detection_results(self, text, image_size): """ 从生成的文本中解析检测结果。 Args: text (str): 模型生成的文本 image_size (tuple): 输入图像的大小(宽度, 高度) Returns: tuple: (boxes, scores, labels) """ # 这是一个简化的示例 - 实际解析将取决于模型输出格式 # 为了演示，我们将提取一些模拟检测结果 # 检查文本中常见的缺陷关键词 defects = [] if "crack" in text.lower() or "裂缝" in text.lower(): defects.append(("裂缝", 0.92, [0.2, 0.3, 0.4, 0.5])) if "corrosion" in text.lower() or "腐蚀" in text.lower(): defects.append(("腐蚀", 0.85, [0.6, 0.2, 0.8, 0.4])) if "damage" in text.lower() or "损坏" in text.lower(): defects.append(("损坏", 0.78, [0.1, 0.7, 0.3, 0.9])) if "defect" in text.lower() or "缺陷" in text.lower(): defects.append(("缺陷", 0.88, [0.5, 0.5, 0.7, 0.7])) # 如果没有找到缺陷，添加一个通用的 if not defects: defects.append(("异常", 0.75, [0.4, 0.4, 0.6, 0.6])) # 将归一化坐标转换为像素坐标 width, height = image_size boxes = [] scores = [] labels = [] for label, score, box in defects: x1, y1, x2, y2 = box pixel_box = [ int(x1 * width), int(y1 * height), int(x2 * width), int(y2 * height) ] boxes.append(pixel_box) scores.append(score) labels.append(label) return boxes, scores, labels