尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

yoloe 提示词检测物体 开放词汇

yoloe 提示词检测物体 开放词汇 目录安装docker部署mobileclip_blt.pt下载封装server 检测移除小框封装server 分割移除小框检测为空bug调查客户端调用封装安装pip install --no-build-isolation openai-whisper pip install githttps://github.com/THU-MIG/yoloe.git#subdirectorythird_party/CLIP pip install githttps://github.com/THU-MIG/yoloe.git#subdirectorythird_party/ml-mobileclip pip install githttps://github.com/THU-MIG/yoloe.git#subdirectorythird_party/lvis-api pip install githttps://github.com/THU-MIG/yoloe.gitdocker部署docker pull ultralytics/ultralytics:latestdocker run --gpus all --ipchost \ -v /data/lbg/project/yoloe/cache:/root/.cache/ultralytics \ -v /data/lbg/project/yoloe/models:/models \ -w /models \ ultralytics/ultralytics:latest \ python -c from ultralytics import YOLOE model YOLOE(yoloe-11m-seg.pt).cuda() model.set_classes([person,car], model.get_text_pe([person,car])) results model.predict(https://ultralytics.com/images/bus.jpg, conf0.25) print(检测到, len(results[0].boxes), 个目标)编译命令不完整docker build -t yoloe:local -f Dockerfile .# Ultralytics YOLO , AGPL-3.0 license # Builds ultralytics/ultralytics:latest image on DockerHub https://hub.docker.com/r/ultralytics/ultralytics # Image is CUDA-optimized for YOLO11 single/multi-GPU training and inference # Start FROM PyTorch image https://hub.docker.com/r/pytorch/pytorch or nvcr.io/nvidia/pytorch:23.03-py3 FROM pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime # Set environment variables # Avoid DDP error MKL_THREADING_LAYERINTEL is incompatible with libgomp.so.1 library https://github.com/pytorch/pytorch/issues/37377 ENV PYTHONUNBUFFERED1 \ PYTHONDONTWRITEBYTECODE1 \ PIP_NO_CACHE_DIR1 \ PIP_BREAK_SYSTEM_PACKAGES1 \ MKL_THREADING_LAYERGNU \ OMP_NUM_THREADS1 # Downloads to user config dir ADD https://github.com/ultralytics/assets/releases/download/v0.0.0/Arial.ttf \ https://github.com/ultralytics/assets/releases/download/v0.0.0/Arial.Unicode.ttf \ /root/.config/Ultralytics/ # Install linux packages # g required to build tflite_support and lap packages, libusb-1.0-0 required for tflite_support package # libsm6 required by libqxcb to create QT-based windows for visualization; set QT_DEBUG_PLUGINS1 to test in docker RUN apt-get update \ apt-get install -y --no-install-recommends \ gcc git zip unzip wget curl htop libgl1 libglib2.0-0 libpython3-dev gnupg g libusb-1.0-0 libsm6 \ rm -rf /var/lib/apt/lists/* # Security updates # https://security.snyk.io/vuln/SNYK-UBUNTU1804-OPENSSL-3314796 RUN apt upgrade --no-install-recommends -y openssl tar # Create working directory WORKDIR /ultralytics # Copy contents and configure git COPY . . #RUN sed -i /^\[http https:\/\/github\.com\/\]/,1d .git/config RUN if [ -f .git/config ]; then sed -i /^\[http https:\/\/github\.com\/\]/,1d .git/config; fi ADD https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n.pt . # Install pip packages RUN python3 -m pip install --upgrade pip wheel # Note -cu12 must be used with tensorrt) #RUN pip install -e .[export] tensorrt-cu12 albumentations1.4.6 comet pycocotools #RUN pip install .[export] tensorrt-cu12 albumentations1.4.6 comet pycocotools RUN pip install ultralytics[export] tensorrt-cu12 albumentations1.4.6 comet pycocotools # Run exports to AutoInstall packages # Edge TPU export fails the first time so is run twice here RUN yolo export modeltmp/yolo11n.pt formatncnn imgsz32 # Requires Python 3.10, bug with paddlepaddle2.5.0 https://github.com/PaddlePaddle/X2Paddle/issues/991 RUN pip install paddlepaddle2.6.0 x2paddle # Fix error: np.bool was a deprecated alias for the builtin bool segmentation error in Tests RUN pip install numpy1.23.5 # Remove extra build files RUN rm -rf tmp /root/.config/Ultralytics/persistent_cache.json # Usage Examples ------------------------------------------------------------------------------------------------------- # Build and Push # tultralytics/ultralytics:latest sudo docker build -f docker/Dockerfile -t $t . sudo docker push $t # Pull and Run with access to all GPUs # tultralytics/ultralytics:latest sudo docker pull $t sudo docker run -it --ipchost --gpus all $t # Pull and Run with access to GPUs 2 and 3 (inside container CUDA devices will appear as 0 and 1) # tultralytics/ultralytics:latest sudo docker pull $t sudo docker run -it --ipchost --gpus device2,3 $t # Pull and Run with local directory access # tultralytics/ultralytics:latest sudo docker pull $t sudo docker run -it --ipchost --gpus all -v $(pwd)/shared/datasets:/datasets $t # Kill all # sudo docker kill $(sudo docker ps -q) # Kill all image-based # sudo docker kill $(sudo docker ps -qa --filter ancestorultralytics/ultralytics:latest) # DockerHub tag update # tultralytics/ultralytics:latest tnewultralytics/ultralytics:v6.2 sudo docker pull $t sudo docker tag $t $tnew sudo docker push $tnew # Clean up # sudo docker system prune -a --volumes # Update Ubuntu drivers # https://www.maketecheasier.com/install-nvidia-drivers-ubuntu/ # DDP test # python -m torch.distributed.run --nproc_per_node 2 --master_port 1 train.py --epochs 3 # GCP VM from Image # docker.io/ultralytics/ultralytics:latestmobileclip_blt.pt下载https://github.com/ultralytics/assets/releases推理代码#!/usr/bin/env python3 import os import sys from pathlib import Path import glob current_dir os.path.dirname(os.path.abspath(__file__)) os.chdir(current_dir) print(current_dir, current_dir) paths [current_dir, current_dir/../] for path in paths: sys.path.insert(0, path) os.environ[PYTHONPATH] (os.environ.get(PYTHONPATH, ) : path).strip(:) import argparse import torch from ultralytics import YOLOE def infer(model_path, image_path, classes, conf0.25, iou0.45, saveNone): model YOLOE(model_path) head model.model.model[-1] print(f加载模型: {model_path}) print(fis_fused{head.is_fused}, nc{head.nc}) # 生成文本嵌入并设置到模型上不调用会触发 fuse 的 set_classes with torch.no_grad(): tpe model.model.get_text_pe(classes) model.model.pe tpe model.model.names_list list(classes) model.model.names {i: n for i, n in enumerate(classes)} head.nc len(classes) # fuseFalse 是关键保持模型开放词汇 results model.predict( image_path, confconf, iouiou, taskdetect, fuseFalse, verboseFalse, ) r results[0] print(f\n图像: {r.path}) print(f检测到 {len(r.boxes)} 个目标) for box in r.boxes: cls_id int(box.cls[0]) name model.model.names.get(cls_id, str(cls_id)) print(f {name}: conf{float(box.conf[0]):.3f}) if save is None: save str(Path(image_path).with_suffix()) f_pred_{_.join(classes)}.jpg r.save(save) print(f结果已保存: {save}) print(f推理后 is_fused{head.is_fused}) return r if __name__ __main__: parser argparse.ArgumentParser(descriptionYOLOE unfused open-vocab inference) # parser.add_argument(--model, defaultyoloe-v8l-seg.pt, help模型路径) parser.add_argument(--model, defaultyoloe-11m-seg.pt, help模型路径) parser.add_argument(--image, default/data/chenna/yoloe-main/img_lbg/generated_image_1783304488.png, help图片路径) # parser.add_argument(--classes, defaultcartoon animal,cartoon dog,cartoon cat,animal,creature, help逗号分隔的类别名例如 stamp,person,car) parser.add_argument(--classes, defaultobject,thing,cartoon animal,animal,creature, help逗号分隔的类别名例如 stamp,person,car) # parser.add_argument(--classes, defaultfront view object,front facing cartoon animal,front view animal,front view creature, help逗号分隔的类别名例如 stamp,person,car) # parser.add_argument(--classes, defaultback view object,back view cartoon animal,back view animal,back view creature, help逗号分隔的类别名例如 stamp,person,car) parser.add_argument(--conf, typefloat, default0.05, help置信度阈值) parser.add_argument(--iou, typefloat, default0.35, helpNMS IoU 阈值) parser.add_argument(--save, defaultz_11_m_back, help保存路径默认自动生成) args parser.parse_args() classes [c.strip() for c in args.classes.split(,) if c.strip()] img_pathsglob.glob(/data/chenna/yoloe-main/img_lbg/*.png) os.makedirs(args.save, exist_okTrue) for img_path in img_paths: img_nameos.path.basename(img_path) save_path args.savef/{img_name} infer(args.model, img_path, classes, args.conf, args.iou, save_path)封装server 检测移除小框# app.py import os import sys from pathlib import Path from typing import List, Optional, Dict, Any import base64 from io import BytesIO import tempfile import glob import sys import os import cv2 from enum import Enum from typing import Tuple, Optional current_dir os.path.dirname(os.path.abspath(__file__)) os.chdir(current_dir) print(current_dir, current_dir) paths [current_dir, current_dir/../] for path in paths: sys.path.insert(0, path) os.environ[PYTHONPATH] (os.environ.get(PYTHONPATH, ) : path).strip(:) from fastapi import FastAPI, File, UploadFile, Form, HTTPException from fastapi.responses import JSONResponse, FileResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field import torch from ultralytics import YOLOE import numpy as np from PIL import Image import uvicorn app FastAPI(titleYOLOE Open-Vocab Detection API, version1.0.0) # 启用CORS app.add_middleware( CORSMiddleware, allow_origins[*], allow_credentialsTrue, allow_methods[*], allow_headers[*], ) # 全局模型缓存 MODEL_CACHE {} DEFAULT_MODEL_PATH yoloe-11m-seg.pt class DetectionResult(BaseModel): class_id: int class_name: str confidence: float bbox: List[float] # [x1, y1, x2, y2] center: List[float] # [cx, cy] width: float height: float class InferenceRequest(BaseModel): classes: List[str] Field(..., description类别列表) conf: float Field(0.25, description置信度阈值, ge0.0, le1.0) iou: float Field(0.45, descriptionNMS IoU阈值, ge0.0, le1.0) model_path: Optional[str] Field(DEFAULT_MODEL_PATH, description模型路径) return_image: bool Field(False, description是否返回带标注的图像) use_contained_filter: bool Field(True, description是否移除被完全包含的框) class InferenceResponse(BaseModel): success: bool detections: List[DetectionResult] num_detections: int image_path: Optional[str] None image_base64: Optional[str] None error: Optional[str] None def filter_contained_boxes(boxes: List[Dict[str, Any]]) - List[Dict[str, Any]]: 移除被其他框完全包含的检测框保留面积最大的框 在 NMS 之后应用作为额外的后处理 if len(boxes) 1: return boxes # 计算面积并排序从大到小 for b in boxes: x1, y1, x2, y2 b[bbox] # 注意这里使用 bbox 而不是 xyxy b[area] (x2 - x1) * (y2 - y1) sorted_boxes sorted(boxes, keylambda x: x[area], reverseTrue) keep [] for b in sorted_boxes: x1, y1, x2, y2 b[bbox] contained False # 检查是否被已保留的某个框包含 for kept in keep: kx1, ky1, kx2, ky2 kept[bbox] # 使用更宽松的包含判断允许小误差 if (kx1 x1 1 and ky1 y1 1 and kx2 x2 - 1 and ky2 y2 - 1): contained True break if not contained: keep.append(b) # 移除临时添加的 area 字段 for b in keep: if area in b: del b[area] return keep def get_model(model_path: str): 获取或加载模型 if model_path not in MODEL_CACHE: try: model YOLOE(model_path) MODEL_CACHE[model_path] model except Exception as e: raise HTTPException(status_code500, detailf模型加载失败: {str(e)}) return MODEL_CACHE[model_path] def run_inference( model_path: str, image_path: str, classes: List[str], conf: float 0.25, iou: float 0.45, return_image: bool False, use_contained_filter: bool True ): 执行推理 try: model get_model(model_path) head model.model.model[-1] # 生成文本嵌入 with torch.no_grad(): tpe model.model.get_text_pe(classes) model.model.pe tpe model.model.names_list list(classes) model.model.names {i: n for i, n in enumerate(classes)} head.nc len(classes) # 执行推理 results model.predict( image_path, confconf, iouiou, taskdetect, fuseFalse, verboseFalse, ) r results[0] detections_o [] if r.boxes is not None: for box in r.boxes: cls_id int(box.cls[0]) name model.model.names.get(cls_id, str(cls_id)) x1, y1, x2, y2 box.xyxy[0].tolist() cx (x1 x2) / 2 cy (y1 y2) / 2 detections_o.append({ class_id: cls_id, class_name: name, confidence: float(box.conf[0]), bbox: [x1, y1, x2, y2], center: [cx, cy], width: x2 - x1, height: y2 - y1 }) # 应用被包含框过滤 if use_contained_filter and detections_o: # 按置信度排序保留高置信度的框 detections_o.sort(keylambda x: x[confidence], reverseTrue) detections filter_contained_boxes(detections_o) print(f{len(detections_o)} 移除被包含框后剩余 {len(detections)} 个目标) # 转换为 Pydantic 模型 detection_results [ DetectionResult(**d) for d in detections ] result_data { success: True, detections: detection_results, num_detections: len(detection_results), image_path: None, image_base64: None, error: None } # 保存带标注的图像 if return_image: with tempfile.NamedTemporaryFile(suffix.jpg, deleteFalse) as tmp: save_path tmp.name r.save(save_path) result_data[image_path] save_path # 转换为base64 with open(save_path, rb) as f: image_data f.read() result_data[image_base64] base64.b64encode(image_data).decode(utf-8) return result_data except Exception as e: return { success: False, detections: [], num_detections: 0, image_path: None, image_base64: None, error: str(e) } app.get(/) async def root(): return { message: YOLOE Open-Vocab Detection API, endpoints: { /detect: POST - 检测图像中的目标, /health: GET - 健康检查 } } app.get(/health) async def health_check(): 健康检查 return { status: healthy, models_loaded: list(MODEL_CACHE.keys()) } app.post(/detect, response_modelInferenceResponse) async def detect_objects( file: UploadFile File(..., description图像文件), classes: str Form(..., description逗号分隔的类别列表), conf: float Form(0.25, description置信度阈值), iou: float Form(0.45, descriptionNMS IoU阈值), model_path: str Form(DEFAULT_MODEL_PATH, description模型路径), return_image: bool Form(False, description是否返回带标注的图像), use_contained_filter: bool Form(True, description是否移除被完全包含的框) ): # 解析类别 class_list [c.strip() for c in classes.split(,) if c.strip()] if not class_list: raise HTTPException(status_code400, detail至少需要一个类别) try: content await file.read() with tempfile.NamedTemporaryFile( suffixPath(file.filename).suffix or .jpg, deleteFalse ) as tmp: tmp.write(content) tmp_path tmp.name # 执行推理 result run_inference( model_pathmodel_path, image_pathtmp_path, classesclass_list, confconf, iouiou, return_imagereturn_image, use_contained_filteruse_contained_filter ) # 清理临时文件 try: os.unlink(tmp_path) except: pass return InferenceResponse(**result) except Exception as e: raise HTTPException(status_code500, detailf处理失败: {str(e)}) app.get(/detect/image/{image_id}) async def get_detection_image(image_id: str): 获取之前检测时保存的带标注图像如果有 return {message: 请使用返回的image_path访问图像} if __name__ __main__: uvicorn.run( jihe_server:app, host0.0.0.0, port7998, reloadFalse, workers1 )封装server 分割移除小框# codingutf-8 import sys import os current_dir os.path.dirname(os.path.abspath(__file__)) os.chdir(current_dir) print(current_dir, current_dir) paths [current_dir, current_dir/../] paths.append(os.path.join(current_dir, src)) for path in paths: sys.path.insert(0, path) os.environ[PYTHONPATH] (os.environ.get(PYTHONPATH, ) : path).strip(:) import threading from typing import List, Dict, Any import io import torch import numpy as np from fastapi import FastAPI, File, UploadFile, Form, HTTPException from PIL import Image from ultralytics import YOLOE import torchvision.ops as ops from fastapi.middleware.cors import CORSMiddleware # ---------------------------------------------------------------------- # 全局模型与锁 # ---------------------------------------------------------------------- MODEL_PATH yoloe-11m-seg.pt # 使用分割模型 model None model_lock threading.Lock() app FastAPI(titleYOLOE Open-Vocabulary Detection Service) app.add_middleware( CORSMiddleware, allow_origins[http://localhost:5173,http://localhost:7999], allow_credentialsTrue, allow_methods[*], allow_headers[*], ) def load_model(path: str) - None: 加载 YOLOE 模型全局 global model model YOLOE(path) print(fModel loaded from {path} csl:{model.model.model[-1].nc}) app.on_event(startup) async def startup_event(): load_model(MODEL_PATH) def apply_nms_all_classes(boxes: List[Dict[str, Any]], iou_threshold: float 0.35) - List[Dict[str, Any]]: 类别无关的 NMS所有框统一处理 if len(boxes) 1: return boxes boxes_tensor torch.tensor([b[xyxy] for b in boxes], dtypetorch.float32) scores_tensor torch.tensor([b[conf] for b in boxes], dtypetorch.float32) class_ids torch.zeros(len(boxes), dtypetorch.int64) keep_indices ops.nms(boxes_tensor, scores_tensor, iou_threshold) return [boxes[i] for i in keep_indices.tolist()] def filter_contained_boxes(boxes: List[Dict[str, Any]]) - List[Dict[str, Any]]: 移除被其他框完全包含的检测框保留面积最大的框 在 NMS 之后应用作为额外的后处理 if len(boxes) 1: return boxes # 计算面积并排序从大到小 for b in boxes: x1, y1, x2, y2 b[xyxy] b[area] (x2 - x1) * (y2 - y1) sorted_boxes sorted(boxes, keylambda x: x[area], reverseTrue) keep [] for b in sorted_boxes: x1, y1, x2, y2 b[xyxy] contained False # 检查是否被已保留的某个框包含 for kept in keep: kx1, ky1, kx2, ky2 kept[xyxy] # 使用更宽松的包含判断允许小误差 if (kx1 x1 1 and ky1 y1 1 and kx2 x2 - 1 and ky2 y2 - 1): contained True break if not contained: keep.append(b) # 移除临时添加的 area 字段 for b in keep: if area in b: del b[area] return keep app.post(/predict) async def predict( file: UploadFile File(..., description待检测的图像文件), classes: str Form(..., description逗号分隔的类别名例如 person,car,dog), conf: float Form(0.05, description置信度阈值), iou: float Form(0.35, descriptionNMS 的 IoU 阈值), use_class_agnostic_nms: bool Form(True, description是否使用类别无关的 NMS所有框统一处理), use_contained_filter: bool Form(True, description是否移除被完全包含的框), ): 检测图像中的目标 - use_class_agnostic_nms: True 时所有检测框被视为同一类别进行 NMS抑制不同类别之间的重叠框 - use_contained_filter: True 时移除被其他框完全包含的检测框保留面积最大的 # 1. 读取图像 try: contents await file.read() img Image.open(io.BytesIO(contents)).convert(RGB) except Exception as e: print(f读取图像失败: {e}) # 使用默认测试图像 img Image.open(/data/lbg/project/hunyuan3d/Hunyuan3D-Part/img_huoshan/img_1784950881_1.jpg).convert(RGB) img_np np.array(img) # 2. 解析类别列表 cls_list [c.strip() for c in classes.split(,) if c.strip()] if not cls_list: raise HTTPException(status_code400, detail未提供任何类别) # 3. 使用锁进行推理 with model_lock: # 清理旧状态 if hasattr(model.model, pe): model.model.pe None if hasattr(model.model, names_list): model.model.names_list None # 设置文本嵌入和类别映射 with torch.no_grad(): text_pe model.model.get_text_pe(cls_list) model.model.pe text_pe model.model.names_list cls_list.copy() model.model.names {i: n for i, n in enumerate(cls_list)} head model.model.model[-1] if hasattr(head, nc): head.nc len(cls_list) if hasattr(head, cached_pe): head.cached_pe None print(f设置的类别数: {len(cls_list)}, head.nc: {head.nc}) # 使用 segment 任务进行推理 results model.predict( img_np, confconf, iouiou, tasksegment, fuseFalse, verboseFalse, ) # 4. 提取检测结果忽略分割掩码 r results[0] boxes_data [] if r.boxes is not None: for box in r.boxes: xyxy box.xyxy[0].tolist() conf_val float(box.conf[0]) cls_id int(box.cls[0]) name model.model.names.get(cls_id, str(cls_id)) boxes_data.append({ xyxy: xyxy, conf: conf_val, cls: cls_id, name: name, }) print(f初始检测到 {len(boxes_data)} 个目标) # 5. 应用后处理 if boxes_data: # 5.1 按置信度排序 boxes_data.sort(keylambda x: x[conf], reverseTrue) # 5.2 应用类别无关的 NMS if use_class_agnostic_nms: boxes_data apply_nms_all_classes(boxes_data, iou_thresholdiou) print(fNMS 后剩余 {len(boxes_data)} 个目标) # 5.3 移除被完全包含的框 if use_contained_filter: boxes_data filter_contained_boxes(boxes_data) print(f移除被包含框后剩余 {len(boxes_data)} 个目标) # 6. 返回 JSON return { boxes: boxes_data, num_detections: len(boxes_data), image_size: [img.width, img.height], nms_type: class_agnostic if use_class_agnostic_nms else class_specific, contained_filter: use_contained_filter, } if __name__ __main__: import uvicorn uvicorn.run(app, host0.0.0.0, port7999)检测为空bug调查有时检测为空原因是项目路径不对没有找到模型路径。客户端调用封装# !/usr/bin/env python3 import cv2 import requests import io import os from PIL import Image, ImageDraw, ImageFont from api_guanfang.img_direct import chat_with_images_stream def predict_sync(image_path, classes, conf0.05, iou0.35, server_urlhttp://localhost:8000): # 打开图像文件 with open(image_path, rb) as f: files {file: (image_path, f, image/jpeg)} data {classes: ,.join(classes), conf: conf, iou: iou, } # 发送 POST 请求 response requests.post(f{server_url}/predict, filesfiles, datadata, ) if response.status_code ! 200: raise Exception(f请求失败: {response.status_code} - {response.text}) result response.json() # 打印结果 print(f图像尺寸: {result[image_size]}) print(f检测到 {result[num_detections]} 个目标) for idx, box in enumerate(result[boxes], 1): x1, y1, x2, y2 box[xyxy] print(fbbox[{x1:.1f}, {y1:.1f}, {x2:.1f}, {y2:.1f}] {idx}. {box[name]}: conf{box[conf]:.3f}, ) return result def expand_box(x1, y1, x2, y2, scale,img_w,img_h): cx (x1 x2) / 2 cy (y1 y2) / 2 w (x2 - x1) * scale h (y2 - y1) * scale new_x1 cx - w / 2 new_y1 cy - h / 2 new_x2 cx w / 2 new_y2 cy h / 2 # 边界限制在图像范围内 new_x1 max(0, int(new_x1)) new_y1 max(0, int(new_y1)) new_x2 min(img_w, int(new_x2)) new_y2 min(img_h, int(new_y2)) return (new_x1, new_y1, new_x2, new_y2) def crop_boxes(image_path, result, expand_scale1.15, avoid_overlapTrue, output_dircropped): os.makedirs(output_dir, exist_okTrue) # 加载原图 img cv2.imread(image_path) if img is None: raise ValueError(f无法读取图片: {image_path}) img_h, img_w img.shape[:2] boxes result[boxes] n len(boxes) if n 0: print(没有检测到目标无需裁剪) return [] # 存储每个框的原始框和放大框 original_boxes [(int(b[xyxy][0]), int(b[xyxy][1]), int(b[xyxy][2]), int(b[xyxy][3])) for b in boxes] expanded_boxes [expand_box(*orig, expand_scale, img_w, img_h) for orig in original_boxes] # 2. 检测重叠两个矩形是否有交集 def rects_overlap(r1, r2): x1, y1, x2, y2 r1 x3, y3, x4, y4 r2 return not (x2 x3 or x4 x1 or y2 y3 or y4 y1) # 3. 避让重叠如果有任何两个放大框重叠则这两个框都回退到原始框 final_boxes expanded_boxes[:] # 先默认全部放大 if avoid_overlap: overlap_pairs [] for i in range(n): for j in range(i 1, n): if rects_overlap(expanded_boxes[i], expanded_boxes[j]): overlap_pairs.append((i, j)) if overlap_pairs: # 收集所有参与重叠的索引 conflict_indices set() for i, j in overlap_pairs: conflict_indices.add(i) conflict_indices.add(j) # 将这些索引对应的框回退到原始 for idx in conflict_indices: final_boxes[idx] original_boxes[idx] print(f检测到 {len(conflict_indices)} 个框不放大) cropped_paths [] for idx, (x1, y1, x2, y2) in enumerate(final_boxes): if x2 x1 or y2 y1: print(f警告: 框 {idx} 无效跳过) continue crop img[y1:y2, x1:x2].copy() # 调用您的 API假设支持 base64 传入 # 注意这里需要根据您的 chat_with_images_stream 接口调整 try: response chat_with_images_stream(text请返回是角色的朝向不要输出过程只要结果, image_paths[crop] # 假设接口支持 base64 ) content response.get(choices, [{}])[0].get(message, {}).get(content, ) except Exception as e: print(f朝向识别失败: {e}) content # 生成文件名 file_name fobj_{idx} if 背 in content: file_name fback_{idx} elif 左侧面 in content: file_name fleft_{idx} elif 右侧面 in content: file_name fright_{idx} elif 正 in content: file_name ffront_{idx} out_path os.path.join(output_dir, f{file_name}.jpg) cv2.imwrite(out_path, crop) cropped_paths.append(out_path) print(f裁剪保存: {out_path} (尺寸 {x2 - x1}x{y2 - y1})) return cropped_paths def visualize_results(image_path, result, output_pathNone): img Image.open(image_path).convert(RGB) draw ImageDraw.Draw(img) # 颜色映射 colors [red, green, blue, yellow, purple, orange, pink, cyan, magenta, lime] for idx, box in enumerate(result[boxes]): x1, y1, x2, y2 box[xyxy] color colors[idx % len(colors)] # 绘制矩形框 draw.rectangle([x1, y1, x2, y2], outlinecolor, width3) # 绘制标签 label f{box[name]}: {box[conf]:.2f} # 获取文本大小简单方法 try: font ImageFont.truetype(/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf, 16) except: font ImageFont.load_default() # 计算文本背景 bbox draw.textbbox((x1, y1), label, fontfont) if hasattr(draw, textbbox) else None if bbox: draw.rectangle(bbox, fillcolor) else: # 旧版 Pillow 兼容 draw.rectangle([x1, y1 - 20, x1 len(label) * 10, y1], fillcolor) draw.text((x1, y1), label, fillwhite, fontfont) # 保存或显示 if output_path is None: output_path image_path.replace(., _detected.) img.save(output_path) print(f可视化结果已保存: {output_path}) return output_path if __name__ __main__: server_url http://192.168.8.231:8000 server_url http://192.168.100.201:7999 image_path rE:\pro_math\math_image\jimeng\image_1783318594.png image_path rE:\pro_math\math_image\jimeng\image_1783319009.png CLASSES [object, thing, cartoon animal, animal, creature] image_path rE:\pro_math\math_image\jimeng\output\source\img_1786419870.jpg CLASSES [bordered square object, bordered circle object, bordered oval object, bordered round-coaster, bordered oval-tin] image_path rC:\Users\ChanJing-01\Pictures\jiezhi\jiezhi2.png CLASSES [ring_simple_band, ring_gemstone, ring_diamond, ring_vintage, ring_modern_art, ring_gold, ring_silver, ring_platinum,watch band,watch bracelet, ring,finger ring,tire,outer tire,casing] try: result predict_sync(image_pathimage_path, classesCLASSES, conf0.05, iou0.2,server_urlserver_url ) cropped_files crop_boxes(image_path, result, expand_scale1.15, avoid_overlapTrue) print(f共裁剪出 {len(cropped_files)} 个目标区域) visualize_results(image_path, result, output_detected.jpg) except Exception as e: print(f错误: {e})
返回列表