yolov12/app.py

# --------------------------------------------------------
# Based on yolov10
# https://github.com/THU-MIG/yolov10/app.py
# --------------------------------------------------------'

import logging
import tempfile
import threading
import cv2
import gradio as gr
import uvicorn
from fastapi import FastAPI
from fastapi import status
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from pydantic import BaseModel

from ultralytics import YOLO
import os
import glob
from typing import Optional

# 设置日志格式和级别
logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s - %(message)s')

def yolov12_inference(image, video, model_id, image_size, conf_threshold):
    model = YOLO(model_id)
    if image:
        results = model.predict(source=image, imgsz=image_size, conf=conf_threshold)
        annotated_image = results[0].plot()
        return annotated_image[:, :, ::-1], None
    else:
        video_path = tempfile.mktemp(suffix=".webm")
        with open(video_path, "wb") as f:
            with open(video, "rb") as g:
                f.write(g.read())

        cap = cv2.VideoCapture(video_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        output_video_path = tempfile.mktemp(suffix=".webm")
        out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'vp80'), fps, (frame_width, frame_height))

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            results = model.predict(source=frame, imgsz=image_size, conf=conf_threshold)
            annotated_frame = results[0].plot()
            out.write(annotated_frame)

        cap.release()
        out.release()

        return None, output_video_path


def yolov12_inference_for_examples(image, model_path, image_size, conf_threshold):
    annotated_image, _ = yolov12_inference(image, None, model_path, image_size, conf_threshold)
    return annotated_image


def app():
    with gr.Blocks():
        with gr.Row():
            with gr.Column():
                image = gr.Image(type="pil", label="Image", visible=True)
                video = gr.Video(label="Video", visible=False)
                input_type = gr.Radio(
                    choices=["Image", "Video"],
                    value="Image",
                    label="Input Type",
                )
                model_id = gr.Dropdown(
                    label="Model",
                    choices=[
                        "yolov12n.pt",
                        "yolov12s.pt",
                        "yolov12m.pt",
                        "yolov12l.pt",
                        "yolov12x.pt",
                    ],
                    value="yolov12m.pt",
                )
                image_size = gr.Slider(
                    label="Image Size",
                    minimum=320,
                    maximum=1280,
                    step=32,
                    value=640,
                )
                conf_threshold = gr.Slider(
                    label="Confidence Threshold",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.05,
                    value=0.25,
                )
                yolov12_infer = gr.Button(value="Detect Objects")

            with gr.Column():
                output_image = gr.Image(type="numpy", label="Annotated Image", visible=True)
                output_video = gr.Video(label="Annotated Video", visible=False)

        def update_visibility(input_type):
            image = gr.update(visible=True) if input_type == "Image" else gr.update(visible=False)
            video = gr.update(visible=False) if input_type == "Image" else gr.update(visible=True)
            output_image = gr.update(visible=True) if input_type == "Image" else gr.update(visible=False)
            output_video = gr.update(visible=False) if input_type == "Image" else gr.update(visible=True)

            return image, video, output_image, output_video

        input_type.change(
            fn=update_visibility,
            inputs=[input_type],
            outputs=[image, video, output_image, output_video],
        )

        def run_inference(image, video, model_id, image_size, conf_threshold, input_type):
            if input_type == "Image":
                return yolov12_inference(image, None, model_id, image_size, conf_threshold)
            else:
                return yolov12_inference(None, video, model_id, image_size, conf_threshold)


        yolov12_infer.click(
            fn=run_inference,
            inputs=[image, video, model_id, image_size, conf_threshold, input_type],
            outputs=[output_image, output_video],
        )

        gr.Examples(
            examples=[
                [
                    "ultralytics/assets/bus.jpg",
                    "yolov12s.pt",
                    640,
                    0.25,
                ],
                [
                    "ultralytics/assets/zidane.jpg",
                    "yolov12x.pt",
                    640,
                    0.25,
                ],
            ],
            fn=yolov12_inference_for_examples,
            inputs=[
                image,
                model_id,
                image_size,
                conf_threshold,
            ],
            outputs=[output_image],
            cache_examples='lazy',
        )

gradio_app = gr.Blocks()
with gradio_app:
    gr.HTML(
        """
    <h1 style='text-align: center'>
    YOLOv12: Attention-Centric Real-Time Object Detectors
    </h1>
    """)
    gr.HTML(
        """
        <h3 style='text-align: center'>
        <a href='https://arxiv.org/abs/2502.12524' target='_blank'>arXiv</a> | <a href='https://github.com/sunsmarterjie/yolov12' target='_blank'>github</a>
        </h3>
        """)
    with gr.Row():
        with gr.Column():
            app()

def start_gradio():
    gradio_app.launch(server_name="0.0.0.0", server_port=7860)

# FastAPI部分
app_fastapi = FastAPI()

class TrainParams(BaseModel):
    """
    用于接收/yolov12/train接口的训练参数，所有参数均需前端传入。
    """
    model: str          # 训练底模
    data: str           # 数据集配置文件路径
    epochs: int         # 训练轮数
    batch: int          # 批次大小
    imgsz: int          # 输入图片尺寸
    scale: float        # 随机缩放增强比例
    mosaic: float       # mosaic数据增强概率
    mixup: float        # mixup数据增强概率
    copy_paste: float   # copy-paste数据增强概率
    device: str         # 训练设备
    project: str        # 工程名
    name: str           # 实验名
    exist_ok: bool      # 是否允许覆盖同名目录

@app_fastapi.post("/yolov12/train")
def yolov12_train(params: TrainParams):
    """
    RESTful POST接口：/yolov12/train
    接收训练参数，调用YOLO模型训练，并返回训练结果。
    返回格式：{"code": 0/1, "msg": "success/错误原因", "result": 训练结果或None}
    """
    logging.info("收到/yolov12/train训练请求")
    logging.info(f"请求参数: {params}")
    try:
        # 根据params.model动态确定配置文件
        if params.model.endswith('.pt'):
            # 如果是.pt文件，将后缀替换为.yaml
            config_file = params.model.replace('.pt', '.yaml')
        else:
            # 如果不是.pt文件，使用默认配置
            config_file = "yolov12.yaml"
        
        model = YOLO(config_file)
        model.load(params.model)
        logging.info("开始模型训练...")
        results = model.train(
            data=params.data,
            epochs=params.epochs,
            batch=params.batch,
            imgsz=params.imgsz,
            scale=params.scale,
            mosaic=params.mosaic,
            mixup=params.mixup,
            copy_paste=params.copy_paste,
            device=params.device,
            project=params.project,
            name=params.name,
            exist_ok=params.exist_ok,
        )
        logging.info("模型训练完成")
        # logging.info(f"训练结果: {str(results)}")
        return {
            "code": 0,
            "msg": "success",
            "result": str(results.save_dir)
        }
    except Exception as e:
        logging.error(f"训练过程发生异常: {e}")
        return {
            "code": 1,
            "msg": str(e),
            "result": None
        }

class PredictParams(BaseModel):
    """
    用于接收/yolov12/predict接口的预测参数，与YOLO predict方法保持一致。
    """
    model: str = "yolov12m.pt"  # 模型路径
    source: str = None  # 输入源（图片/视频路径、URL等）
    stream: bool = False  # 是否流式处理
    conf: float = 0.25  # 置信度阈值
    iou: float = 0.7  # IoU阈值
    max_det: int = 300  # 最大检测数量
    imgsz: int = 640  # 输入图片尺寸
    batch: int = 1  # 批次大小
    device: str = ""  # 设备
    show: bool = False  # 是否显示结果
    save: bool = False  # 是否保存结果
    save_txt: bool = False  # 是否保存txt文件
    save_conf: bool = False  # 是否保存置信度
    save_crop: bool = False  # 是否保存裁剪图片
    show_labels: bool = True  # 是否显示标签
    show_conf: bool = True  # 是否显示置信度
    show_boxes: bool = True  # 是否显示边界框
    line_width: int = None  # 线条宽度
    vid_stride: int = 1  # 视频帧步长
    stream_buffer: bool = False  # 流缓冲区
    visualize: bool = False  # 可视化特征
    augment: bool = False  # 数据增强
    agnostic_nms: bool = False  # 类别无关NMS
    classes: list = None  # 指定类别
    retina_masks: bool = False  # 高分辨率分割掩码
    embed: list = None  # 特征向量层
    half: bool = False  # 半精度
    dnn: bool = False  # OpenCV DNN
    project: str = ""  # 项目名
    name: str = ""  # 实验名
    exist_ok: bool = False  # 是否覆盖现有目录
    verbose: bool = True  # 详细输出

@app_fastapi.post("/yolov12/predict")
def yolov12_predict(params: PredictParams):
    """
    RESTful POST接口：/yolov12/predict
    接收预测参数，调用YOLO模型进行预测，并返回预测结果。
    返回格式：{"code": 0/1, "msg": "success/错误原因", "result": {"save_dir": "保存目录", "filename": "文件名"}}
    """
    logging.info("收到/yolov12/predict预测请求")
    logging.info(f"请求参数: {params}")
    try:
        model = YOLO(params.model)
        logging.info("开始模型预测...")
        
        # 构建预测参数
        predict_kwargs = {}
        for field, value in params.dict().items():
            if field not in ['model'] and value is not None:
                predict_kwargs[field] = value
        
        # 确保保存结果，并强制使用MP4格式
        predict_kwargs['save'] = True
        
        results = model.predict(**predict_kwargs)
        logging.info("模型预测完成")
        
        # 获取保存目录和最终文件名
        result = results[0]
        save_dir = result.save_dir if hasattr(result, 'save_dir') else None
        
        # 获取最终生成的文件名
        final_filename = None
        if save_dir:
            # 获取输入文件名（不含扩展名）
            source = params.source
            base_name = None
            if source:
                base_name = os.path.splitext(os.path.basename(source))[0]
            # 支持的扩展名
            exts = ['*.jpg', '*.jpeg', '*.png', '*.mp4', '*.webm', '*.avi', '*.mov']
            matched_files = []
            for ext in exts:
                matched_files.extend(glob.glob(os.path.join(save_dir, ext)))
            # 按时间排序，查找与输入文件同名的第一个文件
            if base_name and matched_files:
                matched_files = sorted(matched_files, key=os.path.getmtime)
                for f in matched_files:
                    if os.path.splitext(os.path.basename(f))[0] == base_name:
                        final_filename = os.path.basename(f)
                        logging.info(f"按输入文件名查找，返回文件: {final_filename}")
                        break
            # 如果没找到同名文件，返回最新文件
            if not final_filename and matched_files:
                latest_file = max(matched_files, key=os.path.getmtime)
                final_filename = os.path.basename(latest_file)
                logging.info(f"未找到同名，返回最新文件: {final_filename}")
        
        return {
            "code": 0,
            "msg": "success",
            "result": save_dir+"/"+final_filename
        }
    except Exception as e:
        logging.error(f"预测过程发生异常: {e}")
        return {
            "code": 1,
            "msg": str(e),
            "result": None
        }

class StreamParams(BaseModel):
    """
    用于接收 /yolov12/stream 接口的参数。
    model: 推理模型路径
    source: 拉流地址（如rtsp/http视频流）
    stream_url: 推流地址（如rtmp推流地址）
    其他参数同 predict
    """
    model: str = "yolov12m.pt"
    source: str = None
    stream_url: str = None
    conf: float = 0.25
    iou: Optional[float] = 0.7
    imgsz: int = 640
    device: str = ""
    # 可根据需要补充更多参数

@app_fastapi.post("/yolov12/stream")
def yolov12_stream(params: StreamParams):
    """
    RESTful POST接口：/yolov12/stream
    接收视频拉流地址和推流地址，调用YOLO模型推理，将推理后的视频推送到推流地址。
    返回格式：{"code": 0/1, "msg": "success/错误原因", "result": None}
    """
    import cv2
    import logging
    logging.info("收到/yolov12/stream请求")
    logging.info(f"请求参数: {params}")
    try:
        model = YOLO(params.model)
        cap = cv2.VideoCapture(params.source)
        if not cap.isOpened():
            return {"code": 1, "msg": f"无法打开视频流: {params.source}", "result": None}
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        # 推流地址通常为rtmp/rtsp等
        fourcc = cv2.VideoWriter_fourcc(*'flv1') if params.stream_url.startswith('rtmp') else cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(params.stream_url, fourcc, fps if fps > 0 else 25, (width, height))
        if not out.isOpened():
            cap.release()
            return {"code": 1, "msg": f"无法打开推流地址: {params.stream_url}", "result": None}
        frame_count = 0
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            # 推理
            results = model.predict(source=frame, imgsz=params.imgsz, conf=params.conf, iou=params.iou, device=params.device)
            annotated_frame = results[0].plot()
            out.write(annotated_frame)
            frame_count += 1
        cap.release()
        out.release()
        logging.info(f"推理并推流完成，共处理帧数: {frame_count}")
        return {"code": 0, "msg": "success", "result": None}
    except Exception as e:
        logging.error(f"/yolov12/stream 发生异常: {e}")
        return {"code": 1, "msg": str(e), "result": None}

# 全局异常处理器：参数校验失败时统一返回格式
@app_fastapi.exception_handler(RequestValidationError)
async def validation_exception_handler(request, exc):
    err_msg = f"参数校验失败: 路径={request.url.path}, 错误={exc.errors()}"
    logging.error(err_msg)
    return JSONResponse(
        status_code=status.HTTP_200_OK,
        content={
            "code": 422,
            "msg": err_msg,
            "result": None
        }
    )

if __name__ == "__main__":
    threading.Thread(target=start_gradio, daemon=True).start()
    uvicorn.run(app_fastapi, host="0.0.0.0", port=8000)