Spaces:
Sleeping
Sleeping
File size: 3,895 Bytes
e8cd255 f0d0031 e8cd255 f0d0031 e8cd255 f0d0031 e8cd255 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 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='Image detection...')
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='Downloading videos from 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'Video detection, step {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 |