update app

main.py

@@ -10,32 +10,28 @@ from fastapi.responses import FileResponse, RedirectResponse
 from pytorchvideo.models.hub import slowfast_r50
 from ultralytics import YOLO
 
-app = FastAPI()
-
-
-# Create absolute paths based on the app directory
+# --- CONFIGURATION ---
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL_PATH = "best_slowfast_theft.pth"
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 UPLOAD_DIR = os.path.join(BASE_DIR, "uploads")
 OUTPUT_DIR = os.path.join(BASE_DIR, "outputs")
 
-# Ensure they exist (as a backup)
-os.makedirs(UPLOAD_DIR, exist_ok=True)
-os.makedirs(OUTPUT_DIR, exist_ok=True)
-
-# --- CONFIG & GLOBALS ---
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-MODEL_PATH = "best_slowfast_theft.pth"
-UPLOAD_DIR = "uploads"
-OUTPUT_DIR = "outputs"
 CLIP_LEN = 32
 IMG_SIZE = 224
 THEFT_THRESHOLD = 0.6
 
-# In-memory store for job progress (In production, use Redis/Database)
-jobs = {} 
+# Ensure directories exist
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+# In-memory job store
+jobs = {}
+
+app = FastAPI(title="AI Theft Detection System")
 
-# --- LOAD MODELS GLOBALLY (Once) ---
-print(f"Loading models on {DEVICE}...")
+# --- MODEL LOADING ---
+print(f"Loading Models on {DEVICE}...")
 yolo = YOLO("yolov8n.pt")
 
 slowfast_model = slowfast_r50(pretrained=False)
@@ -45,99 +41,132 @@ slowfast_model.blocks[-1].proj = nn.Sequential(
     nn.Linear(in_features, 2)
 )
 
-# Load weights
-ckpt = torch.load(MODEL_PATH, map_location=DEVICE)
-state_dict = ckpt["model"] if "model" in ckpt else ckpt
-slowfast_model.load_state_dict(state_dict)
+if os.path.exists(MODEL_PATH):
+    ckpt = torch.load(MODEL_PATH, map_location=DEVICE)
+    state_dict = ckpt["model"] if "model" in ckpt else ckpt
+    slowfast_model.load_state_dict(state_dict)
+    print("SlowFast weights loaded.")
+else:
+    print(f"Warning: {MODEL_PATH} not found. Running with unitialized weights.")
+
 slowfast_model = slowfast_model.to(DEVICE).eval()
-print("Models loaded successfully.")
 
-# --- HELPER FUNCTIONS ---
+# --- VISUALIZATION HELPERS ---
+
+def draw_corner_rect(img, pt1, pt2, color, thickness, r, d):
+    x1, y1 = pt1
+    x2, y2 = pt2
+    # Top Left
+    cv2.line(img, (x1 + r, y1), (x1 + r + d, y1), color, thickness)
+    cv2.line(img, (x1, y1 + r), (x1, y1 + r + d), color, thickness)
+    cv2.ellipse(img, (x1 + r, y1 + r), (r, r), 180, 0, 90, color, thickness)
+    # Top Right
+    cv2.line(img, (x2 - r, y1), (x2 - r - d, y1), color, thickness)
+    cv2.line(img, (x2, y1 + r), (x2, y1 + r + d), color, thickness)
+    cv2.ellipse(img, (x2 - r, y1 + r), (r, r), 270, 0, 90, color, thickness)
+    # Bottom Left
+    cv2.line(img, (x1 + r, y2), (x1 + r + d, y2), color, thickness)
+    cv2.line(img, (x1, y2 - r), (x1, y2 - r - d), color, thickness)
+    cv2.ellipse(img, (x1 + r, y2 - r), (r, r), 90, 0, 90, color, thickness)
+    # Bottom Right
+    cv2.line(img, (x2 - r, y2), (x2 - r - d, y2), color, thickness)
+    cv2.line(img, (x2, y2 - r), (x2, y2 - r - d), color, thickness)
+    cv2.ellipse(img, (x2 - r, y2 - r), (r, r), 0, 0, 90, color, thickness)
+
+def draw_fancy_overlay(frame, avg_prob, theft_flag, frame_counter):
+    h, w, _ = frame.shape
+    
+    # 1. Semi-transparent Header bar
+    overlay = frame.copy()
+    cv2.rectangle(overlay, (0, 0), (w, 80), (30, 30, 30), -1)
+    cv2.addWeighted(overlay, 0.7, frame, 0.3, 0, frame)
+
+    # 2. Scanning Dot (Pulsing)
+    color_status = (0, 255, 0) if not theft_flag else (0, 0, 255)
+    dot_alpha = (np.sin(frame_counter / 4) + 1) / 2
+    if dot_alpha > 0.4:
+        cv2.circle(frame, (40, 40), 10, color_status, -1)
+    cv2.putText(frame, "AI SURVEILLANCE LIVE", (70, 48), cv2.FONT_HERSHEY_DUPLEX, 0.7, (255, 255, 255), 1)
+
+    # 3. Confidence Meter
+    bar_x, bar_y, bar_w, bar_h = w - 350, 30, 300, 25
+    cv2.rectangle(frame, (bar_x, bar_y), (bar_x + bar_w, bar_y + bar_h), (60, 60, 60), -1)
+    fill_w = int(bar_w * avg_prob)
+    # Color transitions: Green -> Orange -> Red
+    bar_color = (0, 255, 0) if avg_prob < 0.4 else (0, 165, 255) if avg_prob < THEFT_THRESHOLD else (0, 0, 255)
+    cv2.rectangle(frame, (bar_x, bar_y), (bar_x + fill_w, bar_y + bar_h), bar_color, -1)
+    cv2.putText(frame, f"Risk Score: {int(avg_prob*100)}%", (bar_x, bar_y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+
+    # 4. Theft Alert Banner
+    if theft_flag:
+        alert_overlay = frame.copy()
+        cv2.rectangle(alert_overlay, (0, h//2 - 60), (w, h//2 + 60), (0, 0, 200), -1)
+        cv2.addWeighted(alert_overlay, 0.5, frame, 0.5, 0, frame)
+        cv2.putText(frame, "CRITICAL ALERT: THEFT DETECTED", (w//2 - 280, h//2 + 15), 
+                    cv2.FONT_HERSHEY_TRIPLEX, 1.2, (255, 255, 255), 2)
+
+# --- PROCESSING LOGIC ---
 
 def preprocess(frames):
     processed = []
     for frame in frames:
         frame = cv2.resize(frame, (IMG_SIZE, IMG_SIZE))
-        frame = frame[:, :, ::-1]  # BGR to RGB
-        frame = frame / 255.0
+        frame = frame[:, :, ::-1] / 255.0
         processed.append(frame)
-    clip = np.array(processed)
-    clip = np.transpose(clip, (3, 0, 1, 2))  # C,T,H,W
+    clip = np.transpose(np.array(processed), (3, 0, 1, 2))
     return torch.tensor(clip).float().unsqueeze(0)
 
 def process_video_task(job_id: str, input_path: str, output_path: str):
     try:
         cap = cv2.VideoCapture(input_path)
-        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        fps = int(cap.get(cv2.CAP_PROP_FPS))
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-
-        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+        w, h = int(cap.get(3)), int(cap.get(4))
+        fps, total_frames = int(cap.get(5)), int(cap.get(7))
+        out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
 
         frame_buffer = deque(maxlen=CLIP_LEN)
         prediction_buffer = deque(maxlen=10)
-        
         frame_counter = 0
 
-        while True:
+        while cap.isOpened():
             ret, frame = cap.read()
-            if not ret:
-                break
-
+            if not ret: break
             frame_counter += 1
+            
             theft_flag = False
             avg_prob = 0.0
-
             results = yolo(frame, verbose=False)
+
             for r in results:
                 if r.boxes is None: continue
                 for box in r.boxes:
-                    cls = int(box.cls[0])
-                    if cls != 0: continue # Person only
-
+                    if int(box.cls[0]) != 0: continue # Only Person
+                    
                     x1, y1, x2, y2 = map(int, box.xyxy[0])
                     crop = frame[y1:y2, x1:x2]
                     if crop.size == 0: continue
 
                     frame_buffer.append(crop)
-
                     if len(frame_buffer) == CLIP_LEN:
                         clip = preprocess(frame_buffer).to(DEVICE)
-                        # SlowFast inputs
                         inputs = [clip[:, :, ::4, :, :], clip]
                         with torch.no_grad():
-                            outputs = slowfast_model(inputs)
-                            probs = torch.softmax(outputs, dim=1)
-                        
-                        theft_prob = probs[0][1].item()
-                        prediction_buffer.append(theft_prob)
+                            probs = torch.softmax(slowfast_model(inputs), dim=1)
+                        prediction_buffer.append(probs[0][1].item())
                         avg_prob = np.mean(prediction_buffer)
 
-                        if avg_prob > THEFT_THRESHOLD:
-                            theft_flag = True
-
-                    color = (0, 0, 255) if theft_flag else (0, 255, 0)
-                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
-
-            # UI Overlays
-            card_text = f"Class: {'THEFT' if avg_prob > THEFT_THRESHOLD else 'Normal'} | Score: {avg_prob:.2f}"
-            cv2.rectangle(frame, (10, 10), (310, 70), (50, 50, 50), -1)
-            cv2.putText(frame, card_text, (20, 45), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
-            if theft_flag:
-                cv2.putText(frame, "THEFT ALERT", (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 3)
+                    # Determine visual state
+                    active_theft = avg_prob > THEFT_THRESHOLD
+                    color = (0, 0, 255) if active_theft else (0, 255, 0)
+                    draw_corner_rect(frame, (x1, y1), (x2, y2), color, 2, 15, 25)
+                    if active_theft: theft_flag = True
 
+            draw_fancy_overlay(frame, avg_prob, theft_flag, frame_counter)
             out.write(frame)
-            
-            # Update Progress
             jobs[job_id]["progress"] = int((frame_counter / total_frames) * 100)
 
         cap.release()
         out.release()
         jobs[job_id]["status"] = "completed"
-        jobs[job_id]["progress"] = 100
-
     except Exception as e:
         jobs[job_id]["status"] = f"failed: {str(e)}"
 
@@ -148,55 +177,30 @@ async def root():
     return RedirectResponse(url="/docs")
 
 @app.post("/detect")
-async def detect_theft(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
+async def detect(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
     job_id = str(uuid.uuid4())
-    input_filename = f"{job_id}_{file.filename}"
-    input_path = os.path.join(UPLOAD_DIR, input_filename)
+    input_path = os.path.join(UPLOAD_DIR, f"{job_id}_{file.filename}")
     output_path = os.path.join(OUTPUT_DIR, f"result_{job_id}.mp4")
 
-    # Save uploaded file
-    with open(input_path, "wb") as buffer:
-        buffer.write(await file.read())
+    with open(input_path, "wb") as f:
+        f.write(await file.read())
 
-    # Initialize job state
-    jobs[job_id] = {
-        "status": "processing",
-        "progress": 0,
-        "output_path": output_path
-    }
-
-    # Run processing in background
+    jobs[job_id] = {"status": "processing", "progress": 0, "output_path": output_path}
     background_tasks.add_task(process_video_task, job_id, input_path, output_path)
-
-    return {"job_id": job_id, "message": "Video processing started"}
-
+    
+    return {"job_id": job_id, "message": "Video analysis started"}
 
 @app.get("/status/{job_id}")
 async def get_status(job_id: str):
-    if job_id not in jobs:
-        raise HTTPException(status_code=404, detail="Job ID not found")
-    
-    return {
-        "job_id": job_id,
-        "status": jobs[job_id]["status"],
-        "progress": f"{jobs[job_id]['progress']}%"
-    }
-
+    if job_id not in jobs: raise HTTPException(404, "Job not found")
+    return jobs[job_id]
 
 @app.get("/download/{job_id}")
-async def download_video(job_id: str):
-    if job_id not in jobs:
-        raise HTTPException(status_code=404, detail="Job ID not found")
-    
-    if jobs[job_id]["status"] != "completed":
-        raise HTTPException(status_code=400, detail="Video is not processed yet")
-
-    return FileResponse(
-        path=jobs[job_id]["output_path"], 
-        filename=f"annotated_{job_id}.mp4",
-        media_type='video/mp4'
-    )
+async def download(job_id: str):
+    if job_id not in jobs or jobs[job_id]["status"] != "completed":
+        raise HTTPException(400, "File not ready or job not found")
+    return FileResponse(jobs[job_id]["output_path"], filename=f"analyzed_{job_id}.mp4")
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8002)
+    uvicorn.run(app, host="0.0.0.0", port=7860)