Spaces:
Paused
Paused
| # -------------------- SAFETY PATCHES -------------------- | |
| import os | |
| import asyncio | |
| # ✅ Redirect config folders to writable /tmp locations | |
| os.environ["STREAMLIT_HOME"] = "/tmp/.streamlit" | |
| os.environ["STREAMLIT_CONFIG_DIR"] = "/tmp/.streamlit" | |
| os.environ["YOLO_CONFIG_DIR"] = "/tmp/ultralytics" | |
| os.environ["MPLCONFIGDIR"] = "/tmp/mpl" | |
| os.environ["STREAMLIT_WATCH_FILE_SYSTEM"] = "false" # Optional: disable file watcher | |
| # ✅ Create these directories if they don't exist | |
| os.makedirs("/tmp/.streamlit", exist_ok=True) | |
| os.makedirs("/tmp/ultralytics", exist_ok=True) | |
| os.makedirs("/tmp/mpl", exist_ok=True) | |
| # ✅ Fix for "no running event loop" in asyncio | |
| try: | |
| asyncio.get_running_loop() | |
| except RuntimeError: | |
| loop = asyncio.new_event_loop() | |
| asyncio.set_event_loop(loop) | |
| # -------------------- STREAMLIT APP -------------------- | |
| import streamlit as st | |
| import cv2 | |
| import tempfile | |
| import numpy as np | |
| import time | |
| from ultralytics import YOLO | |
| from deep_sort_realtime.deepsort_tracker import DeepSort | |
| from insightface.app import FaceAnalysis | |
| st.set_page_config(layout="wide") | |
| st.title("🎥 Real-Time Gender Detection & Person Tracking") | |
| # ✅ Load models once with cache | |
| def load_models(): | |
| yolo_model = YOLO("yolov8n.pt") | |
| tracker = DeepSort(max_age=30, max_cosine_distance=0.4, nn_budget=None) | |
| face_app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider']) | |
| face_app.prepare(ctx_id=0) | |
| return yolo_model, tracker, face_app | |
| video_file = st.file_uploader("Upload a video", type=["mp4", "avi", "mov"]) | |
| if video_file: | |
| # ✅ Save to temp file | |
| temp_video = tempfile.NamedTemporaryFile(delete=False) | |
| temp_video.write(video_file.read()) | |
| yolo_model, tracker, face_app = load_models() | |
| cap = cv2.VideoCapture(temp_video.name) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| gender_map = {} | |
| total_counts = {"Male": 0, "Female": 0} | |
| gender_colors = {"Male": (255, 255, 0), "Female": (255, 0, 255)} | |
| play = st.checkbox("▶️ Play video", value=True) | |
| FRAME_SKIP = 1 | |
| frame_window = st.empty() | |
| frame_id = 0 | |
| while cap.isOpened() and play: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| frame_id += 1 | |
| if frame_id % FRAME_SKIP != 0: | |
| continue | |
| results = yolo_model(frame)[0] | |
| detections = [] | |
| person_boxes = {} | |
| for box in results.boxes: | |
| cls_id = int(box.cls[0].item()) | |
| label = yolo_model.names[cls_id] | |
| conf = float(box.conf[0].item()) | |
| x1, y1, x2, y2 = map(int, box.xyxy[0].tolist()) | |
| if label == "person": | |
| detections.append(([x1, y1, x2, y2], conf, label)) | |
| person_boxes[(x1, y1, x2, y2)] = None | |
| tracks = tracker.update_tracks(detections, frame=frame) | |
| # Gender detection | |
| for box in person_boxes: | |
| x1, y1, x2, y2 = box | |
| face_crop = frame[y1:y2, x1:x2] | |
| if face_crop.size == 0: | |
| continue | |
| faces = face_app.get(face_crop) | |
| for face in faces: | |
| gender = "Male" if face.sex == 1 else "Female" | |
| person_boxes[box] = gender | |
| total_counts[gender] += 1 | |
| # Draw detections | |
| male_count, female_count = 0, 0 | |
| for track in tracks: | |
| if not track.is_confirmed(): | |
| continue | |
| track_id = track.track_id | |
| label = track.get_det_class() | |
| x1, y1, x2, y2 = map(int, track.to_ltrb()) | |
| gender = None | |
| for box, g in person_boxes.items(): | |
| bx1, by1, bx2, by2 = box | |
| if abs(x1 - bx1) < 20 and abs(y1 - by1) < 20: | |
| gender = g | |
| break | |
| color = gender_colors.get(gender, (200, 200, 200)) | |
| cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2) | |
| label_str = f"ID:{track_id}" | |
| if gender: | |
| label_str += f" ({gender})" | |
| if gender == "Male": | |
| male_count += 1 | |
| else: | |
| female_count += 1 | |
| cv2.putText(frame, label_str, (x1, y1 - 10), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) | |
| # Overlay count | |
| cv2.rectangle(frame, (10, 10), (250, 80), (0, 0, 0), -1) | |
| cv2.putText(frame, f"👨 Male: {male_count}", (20, 35), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.7, gender_colors["Male"], 2) | |
| cv2.putText(frame, f"👩 Female: {female_count}", (20, 65), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.7, gender_colors["Female"], 2) | |
| frame_window.image(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), channels="RGB") | |
| time.sleep(1 / fps) | |
| cap.release() | |
| st.success("✅ Video playback finished.") |