app/detector.py - master - Zolaize/AutoRus_mini

app/detector.py
122 строки4 KB
Zolaize
update app/detector.py
13 дек 2025, 00:28
13 дек 2025, 00:2833a2286
import cv2
from ultralytics import YOLO
import torch
from pathlib import Path
from typing import Union, List, Dict

SELECTED_CLASSES = [2, 3, 5, 7]  # Автомобиль, Мотоцикл, Автобус, Грузовик


def load_model(model_path: str, device: str = None):
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Устройство: {device}")
    print(f"GPU доступна: {torch.cuda.is_available()}")

    model = YOLO(model_path).to(device)
    print("Информация о модели:")
    print(f"Количество классов: {model.model.nc}")
    print(f"Имена классов: {model.names}")
    print(f"ID -> Имя: {dict(enumerate(model.names))}")

    return model, device


def detect_image(model, device: str, image_path: str, output_path: str) -> List[Dict]:

    image = cv2.imread(str(image_path))
    if image is None:
        raise ValueError(f"Не удалось загрузить изображение: {image_path}")

    results = model.predict(image, device=device, verbose=False, conf=0.5)
    annotated_img = results[0].plot()

    # Сохраняем изображение с детекциями
    cv2.imwrite(str(output_path), annotated_img)

    # Формируем список детекций
    detections = []
    for box in results[0].boxes:
        if int(box.cls[0]) in SELECTED_CLASSES:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            bbox = box.xyxy[0].tolist()  # [x1, y1, x2, y2]

            detections.append({
                "class_id": cls_id,
                "confidence": conf,
                "bbox": bbox
            })

    return detections


def detect_video(model, device: str, video_path: str, output_path: str) -> List[Dict]:

    cap = cv2.VideoCapture(str(video_path))

    if not cap.isOpened():
        raise ValueError(f"Не удалось открыть видео: {video_path}")

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    print(f"Размер видео: {width}x{height}, FPS: {fps}, Всего кадров: {total_frames}")

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(str(output_path), fourcc, fps, (width, height))

    frame_count = 0
    all_detections = []

    while True:
        ret, frame = cap.read()
        if not ret or frame is None:
            print(f"Конец видео. Обработано кадров: {frame_count}")
            break

        frame_count += 1
        if frame_count % 60 == 0:  # Вывод каждые 60 кадров
            print(f"Обработан кадр {frame_count}")

        results = model.predict(frame, device=device, verbose=False, conf=0.5)
        annotated_frame = results[0].plot()
        out.write(annotated_frame)

        # Собираем детекции по кадрам
        frame_detections = []
        for box in results[0].boxes:
            if int(box.cls[0]) in SELECTED_CLASSES:
                cls_id = int(box.cls[0])
                conf = float(box.conf[0])
                bbox = box.xyxy[0].tolist()

                frame_detections.append({
                    "frame": frame_count,
                    "class_id": cls_id,
                    "confidence": conf,
                    "bbox": bbox
                })

        all_detections.extend(frame_detections)

    cap.release()
    out.release()

    return all_detections


def process_file(model, device: str, input_path: Union[str, Path], output_path: Union[str, Path]) -> List[Dict]:

    input_path = Path(input_path)
    ext = input_path.suffix.lower()

    if ext in ['.jpg', '.jpeg', '.png']:
        print(f"Обработка изображения: {input_path.name}")
        return detect_image(model, device, str(input_path), str(output_path))
    elif ext in ['.mp4', '.avi', '.mov', '.mkv']:
        print(f"Обработка видео: {input_path.name}")
        return detect_video(model, device, str(input_path), str(output_path))
    else:
        raise ValueError(f"Неподдерживаемый формат файла: {ext}")