update app

main.py

@@ -4,6 +4,7 @@ import cv2
 import torch
 import torch.nn as nn
 import numpy as np
+import datetime
 from collections import deque
 from fastapi import FastAPI, UploadFile, File, BackgroundTasks, HTTPException
 from fastapi.responses import FileResponse, RedirectResponse
@@ -21,11 +22,12 @@ CLIP_LEN = 32
 IMG_SIZE = 224
 THEFT_THRESHOLD = 0.6
 
-# Ensure directories exist
+# Ensure directories exist and are writable
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 # In-memory job store
+# Structure: { job_id: { status, progress, output_path, stop_requested, start_time } }
 jobs = {}
 
 app = FastAPI(title="AI Theft Detection System")
@@ -45,9 +47,9 @@ 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.")
+    print("SlowFast weights loaded successfully.")
 else:
-    print(f"Warning: {MODEL_PATH} not found. Running with unitialized weights.")
+    print(f"Warning: {MODEL_PATH} not found. Running with uninitialized weights.")
 
 slowfast_model = slowfast_model.to(DEVICE).eval()
 
@@ -73,46 +75,48 @@ def draw_corner_rect(img, pt1, pt2, color, thickness, r, d):
     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
+def draw_security_card(frame, avg_prob, theft_flag):
+    # Card Settings
+    card_x1, card_y1 = 30, 30
+    card_w, card_h = 500, 180
+    card_x2, card_y2 = card_x1 + card_w, card_y1 + card_h
+    orange_color = (0, 165, 255) # BGR Orange
+
+    # 1. Draw Semi-Transparent Background
     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)
+    cv2.rectangle(overlay, (card_x1, card_y1), (card_x2, card_y2), (30, 30, 30), -1)
+    cv2.addWeighted(overlay, 0.8, frame, 0.2, 0, frame)
+
+    # 2. Draw Orange Border
+    cv2.rectangle(frame, (card_x1, card_y1), (card_x2, card_y2), orange_color, 2)
+
+    # 3. Text Info
+    now = datetime.datetime.now().strftime("%b %d, %Y, %I:%M:%S %p")
+    status_text = "ALERT: THEFT DETECTED" if theft_flag else "STATUS: NO THEFT"
+    status_color = (0, 0, 255) if theft_flag else (220, 220, 220)
+
+    cv2.putText(frame, "THEFT MONITORING SYSTEM", (card_x1 + 20, card_y1 + 40), 
+                cv2.FONT_HERSHEY_DUPLEX, 0.8, (255, 255, 255), 2)
+    
+    cv2.putText(frame, status_text, (card_x1 + 20, card_y1 + 85), 
+                cv2.FONT_HERSHEY_DUPLEX, 1.0, status_color, 2)
+    
+    cv2.putText(frame, f"TIME: {now}", (card_x1 + 20, card_y1 + 125), 
+                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200, 200, 200), 1)
+    
+    cv2.putText(frame, "THEFT DETECTION: ON", (card_x1 + 20, card_y1 + 160), 
+                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
+
+    # 4. Small Risk Bar inside card
+    bar_full_w = card_w - 40
+    fill_w = int(bar_full_w * avg_prob)
+    cv2.rectangle(frame, (card_x1 + 20, card_y2 - 15), (card_x1 + 20 + bar_full_w, card_y2 - 10), (50, 50, 50), -1)
+    cv2.rectangle(frame, (card_x1 + 20, card_y2 - 15), (card_x1 + 20 + fill_w, card_y2 - 10), orange_color, -1)
 
 # --- PROCESSING LOGIC ---
 
 def preprocess(frames):
-    processed = []
-    for frame in frames:
-        frame = cv2.resize(frame, (IMG_SIZE, IMG_SIZE))
-        frame = frame[:, :, ::-1] / 255.0
-        processed.append(frame)
+    processed = [cv2.resize(f, (IMG_SIZE, IMG_SIZE))[:, :, ::-1] / 255.0 for f in frames]
     clip = np.transpose(np.array(processed), (3, 0, 1, 2))
     return torch.tensor(clip).float().unsqueeze(0)
 
@@ -128,6 +132,11 @@ def process_video_task(job_id: str, input_path: str, output_path: str):
         frame_counter = 0
 
         while cap.isOpened():
+            # Check for Stop Request
+            if jobs.get(job_id, {}).get("stop_requested"):
+                jobs[job_id]["status"] = "stopped"
+                break
+
             ret, frame = cap.read()
             if not ret: break
             frame_counter += 1
@@ -139,7 +148,7 @@ def process_video_task(job_id: str, input_path: str, output_path: str):
             for r in results:
                 if r.boxes is None: continue
                 for box in r.boxes:
-                    if int(box.cls[0]) != 0: continue # Only Person
+                    if int(box.cls[0]) != 0: continue # Person only
                     
                     x1, y1, x2, y2 = map(int, box.xyxy[0])
                     crop = frame[y1:y2, x1:x2]
@@ -148,25 +157,27 @@ def process_video_task(job_id: str, input_path: str, output_path: str):
                     frame_buffer.append(crop)
                     if len(frame_buffer) == CLIP_LEN:
                         clip = preprocess(frame_buffer).to(DEVICE)
-                        inputs = [clip[:, :, ::4, :, :], clip]
                         with torch.no_grad():
-                            probs = torch.softmax(slowfast_model(inputs), dim=1)
+                            probs = torch.softmax(slowfast_model([clip[:, :, ::4, :, :], clip]), dim=1)
                         prediction_buffer.append(probs[0][1].item())
                         avg_prob = np.mean(prediction_buffer)
 
-                    # Determine visual state
-                    active_theft = avg_prob > THEFT_THRESHOLD
-                    color = (0, 0, 255) if active_theft else (0, 255, 0)
+                    is_theft = avg_prob > THEFT_THRESHOLD
+                    color = (0, 0, 255) if is_theft else (0, 255, 0)
                     draw_corner_rect(frame, (x1, y1), (x2, y2), color, 2, 15, 25)
-                    if active_theft: theft_flag = True
+                    if is_theft: theft_flag = True
 
-            draw_fancy_overlay(frame, avg_prob, theft_flag, frame_counter)
+            # Draw the Security Card UI
+            draw_security_card(frame, avg_prob, theft_flag)
+            
             out.write(frame)
             jobs[job_id]["progress"] = int((frame_counter / total_frames) * 100)
 
         cap.release()
         out.release()
-        jobs[job_id]["status"] = "completed"
+        if jobs[job_id]["status"] != "stopped":
+            jobs[job_id]["status"] = "completed"
+
     except Exception as e:
         jobs[job_id]["status"] = f"failed: {str(e)}"
 
@@ -176,29 +187,48 @@ def process_video_task(job_id: str, input_path: str, output_path: str):
 async def root():
     return RedirectResponse(url="/docs")
 
+@app.get("/jobs")
+async def list_jobs():
+    return [{"job_id": jid, "status": data["status"], "progress": f"{data['progress']}%"} for jid, data in jobs.items()]
+
 @app.post("/detect")
 async def detect(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
     job_id = str(uuid.uuid4())
-    input_path = os.path.join(UPLOAD_DIR, f"{job_id}_{file.filename}")
+    input_filename = f"{job_id}_{file.filename}"
+    input_path = os.path.join(UPLOAD_DIR, input_filename)
     output_path = os.path.join(OUTPUT_DIR, f"result_{job_id}.mp4")
 
     with open(input_path, "wb") as f:
         f.write(await file.read())
 
-    jobs[job_id] = {"status": "processing", "progress": 0, "output_path": output_path}
-    background_tasks.add_task(process_video_task, job_id, input_path, output_path)
+    jobs[job_id] = {
+        "status": "processing", 
+        "progress": 0, 
+        "output_path": output_path, 
+        "stop_requested": False,
+        "filename": file.filename
+    }
     
-    return {"job_id": job_id, "message": "Video analysis started"}
+    background_tasks.add_task(process_video_task, job_id, input_path, output_path)
+    return {"job_id": job_id, "message": "Video analysis queued"}
 
 @app.get("/status/{job_id}")
 async def get_status(job_id: str):
-    if job_id not in jobs: raise HTTPException(404, "Job not found")
+    if job_id not in jobs: raise HTTPException(404, "Job ID not found")
     return jobs[job_id]
 
+@app.post("/stop/{job_id}")
+async def stop_job(job_id: str):
+    if job_id not in jobs: raise HTTPException(404, "Job ID not found")
+    if jobs[job_id]["status"] == "processing":
+        jobs[job_id]["stop_requested"] = True
+        return {"message": "Stop signal sent to processing thread."}
+    return {"message": f"Job is already {jobs[job_id]['status']}"}
+
 @app.get("/download/{job_id}")
 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")
+    if job_id not in jobs or jobs[job_id]["status"] not in ["completed", "stopped"]:
+        raise HTTPException(400, "File not ready for download")
     return FileResponse(jobs[job_id]["output_path"], filename=f"analyzed_{job_id}.mp4")
 
 if __name__ == "__main__":