import os
import cv2
import shutil
from ultralytics import YOLO
# --- Configuration ---
input_folder = "/home/philippe/Vidéo/video-sweden/"
output_folder = "videos_avec_animaux/"
os.makedirs(output_folder, exist_ok=True)
model = YOLO("
yolov5su.pt") # Ton modèle personnalisé
CONFIDENCE_THRESHOLD = 0.3
TARGET_CLASSES = ["squirrel", "bird", "deer"]
# --- Parcours récursif ---
for root, dirs, files in os.walk(input_folder):
for filename in files:
if not filename.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
continue
filepath = os.path.join(root, filename)
relative_path = os.path.relpath(filepath, input_folder)
output_subfolder = os.path.join(output_folder, os.path.dirname(relative_path))
os.makedirs(output_subfolder, exist_ok=True)
print(f"🎥 Traitement de : {filepath}")
cap = cv2.VideoCapture(filepath)
if not cap.isOpened():
print(f"⚠️ Impossible d'ouvrir : {filepath}")
continue
fps = int(cap.get(cv2.CAP_PROP_FPS))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
annotated_path = os.path.join(output_subfolder, filename)
out = cv2.VideoWriter(annotated_path, fourcc, fps, (width, height))
animal_detected = False
while True:
ret, frame = cap.read()
if not ret:
break
results = model(frame)[0]
if results.boxes is not None:
for box in results.boxes:
cls_id = int(box.cls.item())
conf = float(box.conf.item())
label = model.names[cls_id].lower()
if label in TARGET_CLASSES and conf >= CONFIDENCE_THRESHOLD:
x1, y1, x2, y2 = map(int, box.xyxy.tolist()[0])
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(frame, f"{label} {conf:.2f}", (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
animal_detected = True
out.write(frame)
cap.release()
out.release()
# --- Déplacement si détection ---
if animal_detected:
moved_path = os.path.join(output_subfolder, "original_" + filename)
shutil.move(filepath, moved_path)
print(f"✅ Animal détecté → vidéo annotée et déplacée : {filename}")
else:
os.remove(annotated_path)
print(f"🚫 Aucun animal détecté → vidéo ignorée : {filename}")