Update utils.py

utils.py

@@ -3,18 +3,25 @@ import glob
 import json
 import urllib.request
 from pathlib import Path
-from typing import List, Dict, Union, Tuple, Optional
+from typing import Tuple
 
-import torch
-import pandas as pd
 import numpy as np
 import cv2
-# import yt_dlp
+import yt_dlp
 import gradio as gr
 from ultralytics import YOLO
 
+import pandas as pd
+import seaborn as sns
+import matplotlib.pyplot as plt
+
+plt.style.use('dark_background')
+plt.rcParams.update({'figure.figsize': (12, 20)})
+plt.rcParams.update({'font.size': 9})
+
 
 YOLO_CLASS_NAMES = json.loads(Path('yolo_classes.json').read_text())
+YOLO_CLASS_NAMES = {int(k): v for k, v in YOLO_CLASS_NAMES.items()}
 
 
 def download_model(model_name: str, models_dir: Path, models: dict) -> str:
@@ -106,4 +113,74 @@ def get_csv_annotate(result_video_path: Path) -> str:
 
     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
+    return result_csv_path
+
+
+def get_matplotlib_fig(csv_annotations_path: str):
+    df = pd.read_csv(csv_annotations_path)
+    df_clean = df.dropna(subset=['class_name'])
+
+    fig, axes = plt.subplots(7, 1, figsize=(10, 20), constrained_layout=True)
+
+    sns.histplot(data=df_clean['conf'], kde=True, ax=axes[0])
+    axes[0].set_title('Распределение уверенности детекций')
+    axes[0].set_xlabel('Уверенность')
+    axes[0].set_ylabel('Количество обнаружений')
+
+    sns.boxplot(data=df_clean, x='class_name', y='conf', ax=axes[1])
+    axes[1].set_title('Распределение уверенности детекции по классам')
+    axes[1].set_xlabel('Класс объекта')
+    axes[1].set_ylabel('Уверенность')
+    axes[1].tick_params(axis='x', labelrotation=45)
+
+    sns.countplot(
+        data=df_clean,
+        x='class_name',
+        hue='class_name',
+        order=df_clean['class_name'].value_counts().index,
+        palette='viridis',
+        legend=False,
+        ax=axes[2],
+        )
+    axes[2].set_title('Количество обнаружений объектов по классам')
+    axes[2].set_xlabel('Класс объекта')
+    axes[2].set_ylabel('Количество')
+
+    face_count_per_frame = df.groupby('frame_number')['box_detected'].sum()
+    axes[3].plot(face_count_per_frame.index, face_count_per_frame.values, marker='o', linestyle='-')
+    axes[3].set_title('Частота обнаружения объектов по кадрам')
+    axes[3].set_xlabel('Номер кадра')
+    axes[3].set_ylabel('Количество обнаруженных объектов')
+
+    face_count_per_frame = df.groupby('frame_sec')['box_detected'].sum()
+    axes[4].plot(face_count_per_frame.index, face_count_per_frame.values, marker='o', linestyle='-')
+    axes[4].set_title('Частота обнаружения объектов по секундам')
+    axes[4].set_xlabel('Время (сек)')
+    axes[4].set_ylabel('Количество обнаруженных объектов')
+
+    sns.scatterplot(
+        data=df_clean,
+        x='frame_sec',
+        y='class_name',
+        hue='class_name',
+        palette='deep',
+        s=50,
+        alpha=0.6,
+        legend=True,
+        ax=axes[5],
+        )
+    axes[5].set_title('Временная шкала обнаружения объектов по классам')
+    axes[5].set_xlabel('Время видео (секунды)')
+    axes[5].set_ylabel('Эмоция')
+    axes[5].grid(True, linestyle='--', alpha=0.7)
+    axes[5].legend(title='Классы объектов', bbox_to_anchor=(1.05, 1), loc='upper left')
+
+    emotion_timeline = df.pivot_table(index='frame_sec', columns='class_name', aggfunc='size', fill_value=0)
+    emotion_timeline.plot(kind='area', stacked=True, ax=axes[6])
+    axes[6].set_title('Динамика обнаружения классов во времени')
+    axes[6].set_xlabel('Время видео (секунды)')
+    axes[6].set_ylabel('Количество детекций')
+    axes[6].grid(True, linestyle='--', alpha=0.7)
+    axes[6].legend(title='Классы объектов', bbox_to_anchor=(1.05, 1), loc='upper left')
+
+    return fig