Spaces:
Sleeping
Sleeping
| # import os | |
| # import cv2 | |
| # import torch | |
| # import numpy as np | |
| # import uuid | |
| # import threading | |
| # import gradio as gr | |
| # from fastapi import FastAPI, UploadFile, File, BackgroundTasks, HTTPException | |
| # from fastapi.responses import FileResponse | |
| # from collections import deque | |
| # from pytorchvideo.models.hub import slowfast_r50 | |
| # from ultralytics import YOLO | |
| # import torch.nn as nn | |
| # # --- SETUP & DIRECTORIES --- | |
| # UPLOAD_DIR = "uploads" | |
| # OUTPUT_DIR = "outputs" | |
| # MODEL_PATH = "models/best_slowfast_theft.pth" | |
| # os.makedirs(UPLOAD_DIR, exist_ok=True) | |
| # os.makedirs(OUTPUT_DIR, exist_ok=True) | |
| # JOBS = {} # Track progress | |
| # # --- MODEL LOADING --- | |
| # DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| # yolo = YOLO("yolov8n.pt") | |
| # def load_slowfast(): | |
| # model = slowfast_r50(pretrained=False) | |
| # in_features = model.blocks[-1].proj.in_features | |
| # model.blocks[-1].proj = nn.Sequential( | |
| # nn.Dropout(p=0.5), | |
| # nn.Linear(in_features, 2) | |
| # ) | |
| # if os.path.exists(MODEL_PATH): | |
| # ckpt = torch.load(MODEL_PATH, map_location=DEVICE) | |
| # model.load_state_dict(ckpt["model"] if "model" in ckpt else ckpt) | |
| # model.to(DEVICE).eval() | |
| # return model | |
| # detector_model = load_slowfast() | |
| # # --- DETECTION LOGIC --- | |
| # def process_video_logic(job_id, input_path, output_path): | |
| # cap = cv2.VideoCapture(input_path) | |
| # total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| # fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
| # w, h = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| # out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)) | |
| # frame_buffer = deque(maxlen=32) | |
| # curr = 0 | |
| # while cap.isOpened(): | |
| # ret, frame = cap.read() | |
| # if not ret: break | |
| # curr += 1 | |
| # JOBS[job_id]["progress"] = int((curr/total_frames)*100) | |
| # # Basic YOLO logic (Simplified for speed) | |
| # results = yolo(frame, verbose=False) | |
| # for r in results: | |
| # for box in r.boxes: | |
| # if int(box.cls[0]) == 0: | |
| # x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| # cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) | |
| # out.write(frame) | |
| # cap.release() | |
| # out.release() | |
| # JOBS[job_id]["status"] = "completed" | |
| # # --- FASTAPI ENDPOINTS --- | |
| # app = FastAPI() | |
| # @app.post("/api/detect") | |
| # async def api_detect(background_tasks: BackgroundTasks, file: UploadFile = File(...)): | |
| # job_id = str(uuid.uuid4()) | |
| # in_p = os.path.join(UPLOAD_DIR, f"{job_id}.mp4") | |
| # out_p = os.path.join(OUTPUT_DIR, f"{job_id}.mp4") | |
| # with open(in_p, "wb") as f: f.write(await file.read()) | |
| # JOBS[job_id] = {"progress": 0, "status": "processing", "file": out_p} | |
| # background_tasks.add_task(process_video_logic, job_id, in_p, out_p) | |
| # return {"job_id": job_id} | |
| # @app.get("/api/progress/{job_id}") | |
| # async def api_progress(job_id: str): | |
| # return JOBS.get(job_id, {"error": "not found"}) | |
| # # --- GRADIO FRONTEND --- | |
| # def web_ui_process(video_input): | |
| # if video_input is None: return None | |
| # job_id = str(uuid.uuid4()) | |
| # out_p = os.path.join(OUTPUT_DIR, f"{job_id}.mp4") | |
| # # Run the processing (Sync for Gradio UI to show progress) | |
| # JOBS[job_id] = {"progress": 0, "status": "processing"} | |
| # process_video_logic(job_id, video_input, out_p) | |
| # return out_p | |
| # with gr.Blocks(title="Theft Detection System") as demo: | |
| # gr.Markdown("# 🛡️ AI Theft Detection System") | |
| # with gr.Row(): | |
| # video_in = gr.Video(label="Upload Video") | |
| # video_out = gr.Video(label="Processed Result") | |
| # btn = gr.Button("Detect Theft") | |
| # btn.click(web_ui_process, inputs=video_in, outputs=video_out) | |
| # # --- MOUNT FASTAPI TO GRADIO --- | |
| # # This allows both to run on the same port on Hugging Face | |
| # app = gr.mount_gradio_app(app, demo, path="/") | |
| # if __name__ == "__main__": | |
| # import uvicorn | |
| # uvicorn.run(app, host="0.0.0.0", port=7860) | |
| import os | |
| import cv2 | |
| import torch | |
| import numpy as np | |
| import uuid | |
| import torch.nn as nn | |
| import gradio as gr | |
| from collections import deque | |
| from pytorchvideo.models.hub import slowfast_r50 | |
| from ultralytics import YOLO | |
| # --- CONFIG & DIRECTORIES --- | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| MODEL_PATH = "models/best_slowfast_theft.pth" | |
| os.makedirs("uploads", exist_ok=True) | |
| os.makedirs("outputs", exist_ok=True) | |
| os.makedirs("models", exist_ok=True) | |
| # --- HEATMAP CLASS --- | |
| class Heatmap: | |
| def __init__(self, h, w, decay=0.92): | |
| self.m = np.zeros((h, w), np.float32) | |
| self.decay = decay | |
| self.h, self.w = h, w | |
| def add(self, bbox, intensity=1.0, poly_mask=None): | |
| x1, y1, x2, y2 = [int(v) for v in bbox] | |
| x1, y1 = max(0, x1), max(0, y1) | |
| x2, y2 = min(self.w-1, x2), min(self.h-1, y2) | |
| if x2 <= x1 or y2 <= y1: return | |
| cx, cy = (x1 + x2) // 2, (y1 + y2) // 2 | |
| rx, ry = max(1, (x2 - x1) // 2), max(1, (y2 - y1) // 2) | |
| blob = np.zeros((self.h, self.w), np.float32) | |
| cv2.ellipse(blob, (cx, cy), (rx, ry), 0, 0, 360, intensity, -1) | |
| blob = cv2.GaussianBlur(blob, (0, 0), rx * 0.6, sigmaY=ry * 0.6) | |
| if poly_mask is not None: blob *= poly_mask | |
| self.m = np.clip(self.m + blob, 0, 10.0) | |
| def step(self): self.m *= self.decay | |
| def overlay(self, frame, alpha=0.45, poly_mask=None): | |
| norm = np.clip(self.m / 10.0, 0, 1) | |
| coloured = cv2.applyColorMap((norm * 255).astype(np.uint8), cv2.COLORMAP_JET) | |
| mask3 = np.stack([(norm > 0.05).astype(np.float32)] * 3, -1) | |
| if poly_mask is not None: mask3 *= np.stack([poly_mask] * 3, -1) | |
| return (coloured * mask3 * alpha + frame * (1 - mask3 * alpha)).astype(np.uint8) | |
| # --- LOAD MODELS --- | |
| print("Loading Models...") | |
| yolo = YOLO("yolov8n.pt") | |
| sf_model = slowfast_r50(pretrained=False) | |
| sf_model.blocks[-1].proj = nn.Sequential(nn.Dropout(p=0.5), nn.Linear(sf_model.blocks[-1].proj.in_features, 2)) | |
| if os.path.exists(MODEL_PATH): | |
| ckpt = torch.load(MODEL_PATH, map_location=DEVICE) | |
| sf_model.load_state_dict(ckpt["model"] if "model" in ckpt else ckpt) | |
| sf_model.to(DEVICE).eval() | |
| # --- CORE LOGIC --- | |
| def process_video(video_path, roi_image): | |
| if not video_path: return None | |
| cap = cv2.VideoCapture(video_path) | |
| w, h = int(cap.get(3)), int(cap.get(4)) | |
| fps = int(cap.get(5)) | |
| output_path = f"outputs/out_{uuid.uuid4()}.mp4" | |
| out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) | |
| heatmap = Heatmap(h, w) | |
| # Process ROI Mask from Sketch | |
| poly_mask = None | |
| if roi_image is not None and 'layers' in roi_image: | |
| # Use the sketch layer to create a mask | |
| mask_layer = roi_image['layers'][0] | |
| mask_layer = cv2.resize(mask_layer, (w, h)) | |
| gray = cv2.cvtColor(mask_layer, cv2.COLOR_BGR2GRAY) | |
| _, poly_mask = cv2.threshold(gray, 10, 1.0, cv2.THRESH_BINARY) | |
| poly_mask = poly_mask.astype(np.float32) | |
| person_buffers = {} | |
| prediction_buffers = {} | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: break | |
| heatmap.step() | |
| results = yolo.track(frame, persist=True, verbose=False, classes=[0]) | |
| global_theft = False | |
| if results[0].boxes.id is not None: | |
| boxes = results[0].boxes.xyxy.cpu().numpy() | |
| ids = results[0].boxes.id.cpu().numpy().astype(int) | |
| for box, track_id in zip(boxes, ids): | |
| x1, y1, x2, y2 = map(int, box) | |
| cx, cy = (x1 + x2) // 2, (y1 + y2) // 2 | |
| # ROI Check | |
| if poly_mask is not None and poly_mask[cy, cx] == 0: continue | |
| heatmap.add(box, poly_mask=poly_mask) | |
| if track_id not in person_buffers: | |
| person_buffers[track_id] = deque(maxlen=32) | |
| prediction_buffers[track_id] = deque(maxlen=10) | |
| crop = frame[y1:y2, x1:x2] | |
| if crop.size > 0: person_buffers[track_id].append(crop) | |
| current_score = 0.0 | |
| if len(person_buffers[track_id]) == 32: | |
| processed = [cv2.resize(f, (224, 224))[:,:,::-1]/255.0 for f in person_buffers[track_id]] | |
| clip = torch.tensor(np.transpose(np.array(processed), (3,0,1,2))).float().unsqueeze(0).to(DEVICE) | |
| with torch.no_grad(): | |
| out_sf = sf_model([clip[:, :, ::4, :, :], clip]) | |
| current_score = torch.softmax(out_sf, dim=1)[0][1].item() | |
| prediction_buffers[track_id].append(current_score) | |
| current_score = np.mean(prediction_buffers[track_id]) | |
| if current_score > 0.6: global_theft = True | |
| color = (0, 0, 255) if current_score > 0.6 else (0, 255, 0) | |
| cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2) | |
| cv2.putText(frame, f"ID:{track_id} {current_score:.2f}", (x1, y1-10), 0, 0.5, color, 2) | |
| # Fancy Overlays | |
| frame = heatmap.overlay(frame, poly_mask=poly_mask) | |
| overlay = frame.copy() | |
| cv2.rectangle(overlay, (0,0), (w, 80), (0,0,0), -1) | |
| cv2.addWeighted(overlay, 0.5, frame, 0.5, 0, frame) | |
| status = "!!! THEFT DETECTED !!!" if global_theft else "Monitoring Area..." | |
| scolor = (0,0,255) if global_theft else (0,255,0) | |
| cv2.putText(frame, status, (20, 50), 0, 1.0, scolor, 3) | |
| out.write(frame) | |
| cap.release() | |
| out.release() | |
| return output_path | |
| # --- GRADIO UI --- | |
| with gr.Blocks(theme=gr.themes.Soft(), title="Theft Detection Pro") as demo: | |
| gr.Markdown("# 🛡️ AI Theft Detection & Heatmap System") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| video_input = gr.Video(label="1. Upload Video", height=400) | |
| gr.Markdown("### 2. Draw ROI (Optional)\nDraw on the image below to monitor a specific area.") | |
| # This handles getting the first frame automatically | |
| roi_input = gr.ImageMask(label="Draw Region of Interest", height=400) | |
| def get_first_frame(vid): | |
| if vid is None: return None | |
| cap = cv2.VideoCapture(vid) | |
| ret, frame = cap.read() | |
| cap.release() | |
| if ret: return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| return None | |
| video_input.change(get_first_frame, inputs=video_input, outputs=roi_input) | |
| submit_btn = gr.Button("🚀 Start Processing", variant="primary") | |
| with gr.Column(scale=1): | |
| video_output = gr.Video(label="3. Detection Result", height=800) | |
| submit_btn.click( | |
| fn=process_video, | |
| inputs=[video_input, roi_input], | |
| outputs=video_output | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |