Upload 4 files

app.py

@@ -0,0 +1,183 @@
+import shutil
+from pathlib import Path
+from typing import List, Dict, Union, Tuple, Literal, Optional
+
+import numpy as np
+import gradio as gr
+from gradio.components.base import Component
+from ultralytics import YOLO
+
+from utils import download_model, detect_image, detect_video, get_csv_annotate
+
+
+# ======================= МОДЕЛЬ ===================================
+
+MODELS_DIR = Path('models')
+MODELS_DIR.mkdir(exist_ok=True)
+
+MODELS = {
+    'yolov8n.pt': 'https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8n.pt',
+    'yolov8s.pt': 'https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8s.pt',
+    'yolov8m.pt': 'https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8m.pt',
+    'yolov8l.pt': 'https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8l.pt',
+    'yolov8x.pt': 'https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8x.pt',
+}
+MODEL_NAMES = list(MODELS.keys())
+
+model_path = download_model(MODEL_NAMES[0], MODELS_DIR, MODELS)
+default_model = YOLO(model_path)
+
+IMAGE_EXTENSIONS = ['.jpg', '.jpeg', '.png']
+VIDEO_EXTENSIONS = ['.mp4', '.avi']
+
+
+# =================== ДОП ФУНКЦИИ ИНТРЕФЕЙСА ==============================
+
+def change_model(model_state: Dict[str, YOLO], model_name: str):
+    progress = gr.Progress()
+    progress(0.3, desc='Загрузка модели')
+    model_path = download_model(model_name)
+    progress(0.7, desc='Инициализация модели')
+    model_state['model'] = YOLO(model_path)
+    return f"Модель {model_name} инициализирована"
+
+
+def detect(file_path: str, file_link: str, model_state: Dict[str, YOLO], conf: float, iou: float):
+    model = model_state['model']
+    if file_link:
+        file_path = file_link
+
+    file_ext = f'.{file_path.rsplit(".")[-1]}'
+    if file_ext in IMAGE_EXTENSIONS:
+        np_image = detect_image(file_path, model, conf, iou)
+        return np_image, "Детекция завершена, открытие изображения..."
+    elif file_ext in VIDEO_EXTENSIONS or 'youtube.com' in file_link:
+        video_path = detect_video(file_path, model, conf, iou)
+        return video_path, "Детекция завершена, конвертация и открытие видео..."
+    else:
+        gr.Info('Неверный формат изображения или видео...')
+        return None, None
+
+# =================== КОМПОНЕНТЫ ИНТРЕФЕЙСА ==============================
+
+def get_output_media_components(detect_result: Optional[Union[np.ndarray, str, Path]] = None):
+    visible = isinstance(detect_result, np.ndarray)
+    image_output = gr.Image(
+        value=detect_result if visible else None,
+        type="numpy",
+        width=640,
+        height=480,
+        visible=visible,
+        label='Output',
+        )
+    visible = isinstance(detect_result, (str, Path))
+    video_output = gr.Video(
+        value=detect_result if visible else None,
+        width=640,
+        height=480,
+        visible=visible,
+        label='Output',
+        )
+    clear_btn = gr.Button(
+        value='Clear',
+        scale=0,
+        visible=detect_result is not None,
+        )
+    return image_output, video_output, clear_btn
+
+
+def get_download_csv_btn(csv_annotations_path: Optional[Path] = None):
+    download_csv_btn = gr.DownloadButton(
+        label='Скачать csv аннотации к видео',
+        value=csv_annotations_path,
+        scale=0,
+        visible=csv_annotations_path is not None,
+        )
+    return download_csv_btn
+
+# =================== ИНТЕРФЕЙС ПРИЛОЖЕНИЯ ==========================
+
+css = '''
+.gradio-container { width: 70% !important }
+'''
+with gr.Blocks(css=css) as demo:
+    gr.HTML("""<h3 style='text-align: center'>YOLOv8 Detector</h3>""")
+    
+    model_state = gr.State({'model': default_model})
+    detect_result = gr.State(None)
+    csv_annotations_path = gr.State(None)
+
+    with gr.Row():
+        with gr.Column():
+            file_path = gr.File(file_types=['image', 'video'], file_count='single', label='Выберите изображение или видео')
+            file_link = gr.Textbox(label='Прямая ссылка на изображение или ссылка на YouTube')
+            model_name = gr.Radio(choices=MODEL_NAMES, value=MODEL_NAMES[0], label='Модель YOLO')
+            conf = gr.Slider(0, 1, value=0.5, step=0.05, label='Порог уверенности')
+            iou = gr.Slider(0, 1, value=0.7, step=0.1, label='Порог IOU')
+            status_message = gr.Textbox(value='Готово к работе', label='Статус')
+            detect_btn = gr.Button('Detect', interactive=True)
+
+        with gr.Column():
+            image_output, video_output, clear_btn = get_output_media_components()
+            download_csv_btn = get_download_csv_btn()
+
+    model_name.change(
+        fn=lambda: gr.update(interactive=False),
+        inputs=None,
+        outputs=[detect_btn],
+    ).then(
+        fn=change_model,
+        inputs=[model_state, model_name],
+        outputs=[status_message],
+    ).success(
+        fn=lambda: gr.update(interactive=True),
+        inputs=None,
+        outputs=[detect_btn],
+    )
+
+    detect_btn.click(
+        fn=detect,
+        inputs=[file_path, file_link, model_state, conf, iou],
+        outputs=[detect_result, status_message],
+    ).success(
+        fn=get_output_media_components,
+        inputs=[detect_result],
+        outputs=[image_output, video_output, clear_btn],
+    ).then(
+        fn=lambda: 'Готово к работе',
+        inputs=None,
+        outputs=[status_message],
+    ).then(
+        fn=get_csv_annotate,
+        inputs=[detect_result],
+        outputs=[csv_annotations_path],
+    ).success(
+        fn=get_download_csv_btn,
+        inputs=[csv_annotations_path],
+        outputs=[download_csv_btn],
+    )
+
+    def clear_results_dir(detect_result):
+        if isinstance(detect_result, Path):
+            shutil.rmtree(detect_result.parent, ignore_errors=True)
+
+    clear_components = [image_output, video_output, clear_btn, download_csv_btn]
+    clear_btn.click(
+        fn=lambda: [gr.update(visible=False) for _ in range(len(clear_components))],
+        inputs=None,
+        outputs=clear_components,
+    ).then(
+        fn=clear_results_dir,
+        inputs=[detect_result],
+        outputs=None,
+    ).then(
+        fn=lambda: (None, None),
+        inputs=None,
+        outputs=[detect_result, csv_annotations_path]
+        )
+
+    gr.HTML("""<h3 style='text-align: center'>
+    <a href="https://github.com/sergey21000/yolo_gradio_detector" target='_blank'>GitHub Page</a></h3>
+    """)
+    
+demo.launch(server_name='0.0.0.0')  # debug=True

requirements.txt

@@ -0,0 +1,4 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+ultralytics>=8,<9
+gradio>4
+yt_dlp

utils.py

@@ -0,0 +1,109 @@
+import os
+import glob
+import json
+import urllib.request
+from pathlib import Path
+from typing import List, Dict, Union, Tuple, Optional
+
+import torch
+import pandas as pd
+import numpy as np
+import cv2
+import yt_dlp
+import gradio as gr
+from ultralytics import YOLO
+
+
+YOLO_CLASS_NAMES = json.loads(Path('yolo_classes.json').read_text())
+
+
+def download_model(model_name: str, models_dir: Path, models: dict) -> str:
+    model_path = models_dir / model_name
+    if not model_path.exists():
+        urllib.request.urlretrieve(models[model_name], model_path)
+    return str(model_path)
+
+
+def detect_image(image_path: str, model: YOLO, conf: float, iou: float) -> np.ndarray:
+    gr.Progress()(0.5, desc='Детекция изображения...')
+    detections = model.predict(source=image_path, conf=conf, iou=iou)
+    np_image = detections[0].plot()
+    np_image = cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
+    return np_image
+
+
+def detect_video(video_path_or_url: str, model: YOLO, conf: float, iou: float) -> Tuple[Path, Path]:
+    progress = gr.Progress()
+    video_path = video_path_or_url
+    if 'youtube.com' in video_path_or_url or 'youtu.be' in video_path_or_url:
+        progress(0.001, desc='Загрузка видео с YouTube...')
+        ydl_opts = {'format': 'bestvideo[height<=720]'}
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            video_info_dict = ydl.extract_info(video_path_or_url, download=True)
+            video_path = ydl.prepare_filename(video_info_dict)
+
+    cap = cv2.VideoCapture(video_path)
+    num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    cap.release()
+
+    generator = model.predict(
+        source=video_path,
+        conf=0.5,
+        iou=0.5,
+        save=True,
+        save_txt=True,
+        save_conf=True,
+        stream=True,
+        verbose=False,
+        )
+
+    frames_count = 0
+    for result in generator:
+        frames_count += 1
+        progress((frames_count, num_frames), desc=f'Детекция видео, шаг {frames_count}/{num_frames}')
+
+    file_name = Path(result.path).with_suffix('.avi').name
+    result_video_path = Path(result.save_dir) / file_name
+    Path(video_path).unlink(missing_ok=True)
+    return result_video_path
+
+
+def get_csv_annotate(result_video_path: Path) -> str:
+    if not isinstance(result_video_path, Path):
+        return None
+
+    txts_path = result_video_path.parent / 'labels'
+    escaped_pattern = glob.escape(result_video_path.stem)
+    matching_txts_path = sorted(txts_path.glob(f'{escaped_pattern}_*.txt'), key=os.path.getmtime)
+
+    df_list = []
+    for txt_path in matching_txts_path:
+        frame_number = int(txt_path.stem.rsplit('_')[-1])
+        with open(txt_path) as file:
+            df_rows = file.readlines()
+            for df_row in df_rows:
+                df_row = map(float, df_row.split())
+                df_list.append((frame_number, *df_row))
+
+    column_names = ['frame_number', 'class_label', 'x', 'y', 'w', 'h', 'conf']
+    df = pd.DataFrame(df_list, columns=column_names)
+
+    df.class_label = df.class_label.astype(int)
+    class_name_series = df.class_label.map(YOLO_CLASS_NAMES)
+    df.insert(loc=1, column='class_name', value=class_name_series)
+
+    cap = cv2.VideoCapture(str(result_video_path))
+    frames_fps = int(cap.get(cv2.CAP_PROP_FPS))
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    cap.release()
+    
+    frame_sec_series = df.frame_number / frames_fps
+    df.insert(loc=1, column='frame_sec', value=frame_sec_series)
+
+    full_frames = pd.DataFrame({'frame_number': range(total_frames)})
+    df = pd.merge(full_frames, df, on='frame_number', how='outer')
+    df.frame_sec = df.frame_number / frames_fps
+
+    result_csv_path = f'{result_video_path.parent / result_video_path.stem}_annotations.csv'
+    df.to_csv(result_csv_path, index=False)
+    return result_csv_path

yolo_classes.json

@@ -0,0 +1,82 @@
+{
+    "0": "person",
+    "1": "bicycle",
+    "2": "car",
+    "3": "motorcycle",
+    "4": "airplane",
+    "5": "bus",
+    "6": "train",
+    "7": "truck",
+    "8": "boat",
+    "9": "traffic light",
+    "10": "fire hydrant",
+    "11": "stop sign",
+    "12": "parking meter",
+    "13": "bench",
+    "14": "bird",
+    "15": "cat",
+    "16": "dog",
+    "17": "horse",
+    "18": "sheep",
+    "19": "cow",
+    "20": "elephant",
+    "21": "bear",
+    "22": "zebra",
+    "23": "giraffe",
+    "24": "backpack",
+    "25": "umbrella",
+    "26": "handbag",
+    "27": "tie",
+    "28": "suitcase",
+    "29": "frisbee",
+    "30": "skis",
+    "31": "snowboard",
+    "32": "sports ball",
+    "33": "kite",
+    "34": "baseball bat",
+    "35": "baseball glove",
+    "36": "skateboard",
+    "37": "surfboard",
+    "38": "tennis racket",
+    "39": "bottle",
+    "40": "wine glass",
+    "41": "cup",
+    "42": "fork",
+    "43": "knife",
+    "44": "spoon",
+    "45": "bowl",
+    "46": "banana",
+    "47": "apple",
+    "48": "sandwich",
+    "49": "orange",
+    "50": "broccoli",
+    "51": "carrot",
+    "52": "hot dog",
+    "53": "pizza",
+    "54": "donut",
+    "55": "cake",
+    "56": "chair",
+    "57": "couch",
+    "58": "potted plant",
+    "59": "bed",
+    "60": "dining table",
+    "61": "toilet",
+    "62": "tv",
+    "63": "laptop",
+    "64": "mouse",
+    "65": "remote",
+    "66": "keyboard",
+    "67": "cell phone",
+    "68": "microwave",
+    "69": "oven",
+    "70": "toaster",
+    "71": "sink",
+    "72": "refrigerator",
+    "73": "book",
+    "74": "clock",
+    "75": "vase",
+    "76": "scissors",
+    "77": "teddy bear",
+    "78": "hair drier",
+    "79": "toothbrush"
+}