# 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()