import os import cv2 import os.path as osp from mmengine.config import Config from mmengine.dataset import Compose from mmdet.apis import init_detector from mmdet.utils import get_test_pipeline_cfg # from mmengine.runner.amp import autocast from torch.amp import autocast import torch import supervision as sv from typing import Dict, Optional, Sequence, List import supervision as sv class LabelAnnotator(sv.LabelAnnotator): @staticmethod def resolve_text_background_xyxy( center_coordinates, text_wh, position, ): center_x, center_y = center_coordinates text_w, text_h = text_wh return center_x, center_y, center_x + text_w, center_y + text_h class YoloInterface: def __init__(self): """ Initialize the YOLO-World model with the given configuration and checkpoint. Args: """ pass def set_BBoxAnnotator(self): self.BOUNDING_BOX_ANNOTATOR = sv.BoundingBoxAnnotator(thickness=1) # MASK_ANNOTATOR = sv.MaskAnnotator() self.LABEL_ANNOTATOR = LabelAnnotator(text_padding=4, text_scale=0.5, text_thickness=1) class YoloWorldInterface(YoloInterface): def __init__(self, config_path: str, checkpoint_path: str, device: str = "cuda:0"): """ Initialize the YOLO-World model with the given configuration and checkpoint. Args: config_path (str): Path to the model configuration file. checkpoint_path (str): Path to the model checkpoint. device (str): Device to run the model on (e.g., 'cuda:0', 'cpu'). """ self.config_path = config_path self.checkpoint_path = checkpoint_path self.device = device # Load configuration cfg = Config.fromfile(config_path) cfg.work_dir = osp.join('./work_dirs', osp.splitext(osp.basename(config_path))[0]) cfg.load_from = checkpoint_path # Initialize the model self.model = init_detector(cfg, checkpoint=checkpoint_path, device=device) self.set_BBoxAnnotator() # Initialize the test pipeline # build test pipeline self.model.cfg.test_dataloader.dataset.pipeline[ 0].type = 'mmdet.LoadImageFromNDArray' self.test_pipeline = Compose(self.model.cfg.test_dataloader.dataset.pipeline) def reparameterize_object_list(self, target_objects: List[str], cue_objects: List[str]): """ Reparameterize the detect object list to be used by the YOLO model. Args: target_objects (List[str]): List of target object names. cue_objects (List[str]): List of cue object names. """ # Combine target objects and cue objects into the final text format combined_texts = target_objects + cue_objects # Format the text prompts for the YOLO model self.texts = [[obj.strip()] for obj in combined_texts] + [[' ']] # Reparameterize the YOLO model with the provided text prompts self.model.reparameterize(self.texts) def inference(self, image: str, max_dets: int = 100, score_threshold: float = 0.3, use_amp: bool = False): """ Run inference on a single image. Args: image (str): Path to the image. max_dets (int): Maximum number of detections to keep. score_threshold (float): Score threshold for filtering detections. use_amp (bool): Whether to use mixed precision for inference. Returns: sv.Detections: Detection results. """ # Prepare data for inference data_info = dict(img_id=0, img_path=image, texts=self.texts) data_info = self.test_pipeline(data_info) data_batch = dict(inputs=data_info['inputs'].unsqueeze(0), data_samples=[data_info['data_samples']]) # Run inference with autocast(enabled=use_amp), torch.no_grad(): output = self.model.test_step(data_batch)[0] pred_instances = output.pred_instances pred_instances = pred_instances[pred_instances.scores.float() > score_threshold] if len(pred_instances.scores) > max_dets: indices = pred_instances.scores.float().topk(max_dets)[1] pred_instances = pred_instances[indices] pred_instances = pred_instances.cpu().numpy() # Process detections detections = sv.Detections( xyxy=pred_instances['bboxes'], class_id=pred_instances['labels'], confidence=pred_instances['scores'], mask=pred_instances.get('masks', None) ) return detections def inference_detector(self, images, max_dets=50, score_threshold=0.2, use_amp: bool = False): data_info = dict(img_id=0, img=images[0], texts=self.texts) #TBD for batch searching data_info = self.test_pipeline(data_info) data_batch = dict(inputs=data_info['inputs'].unsqueeze(0), data_samples=[data_info['data_samples']]) detections_inbatch = [] with torch.no_grad(): outputs = self.model.test_step(data_batch) # cover to searcher interface format for output in outputs: pred_instances = output.pred_instances pred_instances = pred_instances[pred_instances.scores.float() > score_threshold] if len(pred_instances.scores) > max_dets: indices = pred_instances.scores.float().topk(max_dets)[1] pred_instances = pred_instances[indices] output.pred_instances = pred_instances if 'masks' in pred_instances: masks = pred_instances['masks'] else: masks = None pred_instances = pred_instances.cpu().numpy() detections = sv.Detections(xyxy=pred_instances['bboxes'], class_id=pred_instances['labels'], confidence=pred_instances['scores'], mask=masks) detections_inbatch.append(detections) self.detect_outputs_raw = outputs self.detections_inbatch = detections_inbatch return detections_inbatch def bbox_visualization(self, images, detections_inbatch): anno_images = [] # detections_inbatch = self.detections_inbatch for b, detections in enumerate(detections_inbatch): texts = self.texts labels = [ f"{texts[class_id][0]} {confidence:0.2f}" for class_id, confidence in zip(detections.class_id, detections.confidence) ] index = len(detections_inbatch) -1 image = images[index] anno_image = image.copy() anno_image = self.BOUNDING_BOX_ANNOTATOR.annotate(anno_image, detections) anno_image = self.LABEL_ANNOTATOR.annotate(anno_image, detections, labels=labels) anno_images.append(anno_image) return anno_images import torch from typing import List import supervision as sv # 确保已安装 Supervision 库 import os.path as osp class YoloV5Interface(YoloInterface): def __init__(self,config_path="ultralytics/yolov5", checkpoint_path: str = 'yolov5s', device: str = 'cuda:0'): """ 初始化 YOLOv5 模型。 Args: model_name (str): YOLOv5 模型变体名称（如 'yolov5s', 'yolov5m', 'yolov5l', 'yolov5x'）。 device (str): 运行模型的设备（如 'cuda:0', 'cpu'）。 """ self.device = device self.model = torch.hub.load("ultralytics/yolov5", "yolov5s", pretrained=True) self.model.to(self.device) self.model.eval() self.target_classes = None # 用于存储目标类别列表 self.texts = None self.test_pipeline = None def reparameterize_object_list(self, target_objects: List[str], cue_objects: List[str]): """ 重新参数化检测对象列表，以便在推理时使用。 Args: target_objects (List[str]): 目标对象名称列表。 cue_objects (List[str]): 线索对象名称列表。 """ # 合并目标对象和线索对象 combined_objects = target_objects + cue_objects self.target_classes = combined_objects def inference(self, images: str, max_dets: int = 100, score_threshold: float = 0.3, use_amp: bool = False): """ 对单张图像运行推理。 Args: image (str): 图像路径。 max_dets (int): 保留的最大检测数量。 score_threshold (float): 过滤检测的分数阈值。 use_amp (bool): 是否使用混合精度进行推理。 Returns: sv.Detections: 检测结果。 """ results = self.model(images, size=640) # 可以根据需要调整输入尺寸 # 提取检测结果（假设批量大小为 1） detections_batch = results.pred # B tensors of shape (N, 6) [x1, y1, x2, y2, confidence, class] # 应用分数阈值 # 用于存储每个批次过滤后的检测结果 filtered_detections = [] for detections in detections_batch: # 应用分数阈值，过滤掉 confidence <= score_threshold 的检测 detections = detections[detections[:, 4] > score_threshold] # 如果设置了 topk，截取前 topk 个检测 if len(detections) > max_dets: detections = detections[:max_dets] # 如果设置了目标类别，过滤检测结果 if self.target_classes is not None: # 获取所有类别名称 class_names = self.model.names # 获取目标类别的类别ID target_class_ids = [i for i, name in class_names.items() if name in self.target_classes] # 过滤检测结果 detections = detections[[cls in target_class_ids for cls in detections[:, 5]]] # 转换为 Supervision 库的 Detections 对象 detections = sv.Detections( xyxy=detections[:, :4].cpu().numpy(), confidence=detections[:, 4].cpu().numpy(), class_id=detections[:, 5].cpu().numpy().astype(int) ) filtered_detections.append(detections) return filtered_detections