Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 import os
 import cv2
 import numpy as np
@@ -8,39 +9,34 @@ from googleapiclient.discovery import build
 from googleapiclient.http import MediaIoBaseDownload, MediaFileUpload
 from io import BytesIO
 
-
 # === CONFIGURATION ===
 CCTVFEED_IDS = ['1KJRkSf2SKEZ1mXS9_si65IwMBtjs6p4n'] # Replace with actual folder IDs
 REG_FOLDER_ID = '1qkcR7nQTEtiMH9OFUv2bGxVn08E3dKjF'
 INTRUDER_FOLDER_ID = '1PPAUWU-wMx7fek73p-hqPqYQypYtG8Ob' 
-TOKEN_PICKLE = 'token.pickle'                        # Must be uploaded to your HF Space
+TOKEN_PICKLE = 'token.pickle'
 CASCADE_PATH = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
+MATCH_THRESHOLD = 8000  # More lenient threshold
 
-# === GOOGLE DRIVE AUTH ===
 def init_gdrive():
     with open(TOKEN_PICKLE, 'rb') as token:
         creds = pickle.load(token)
     return build('drive', 'v3', credentials=creds)
 
-# === DOWNLOAD FILE FROM DRIVE ===
 def download_file(service, file_id, local_path):
     request = service.files().get_media(fileId=file_id)
     fh = BytesIO()
     downloader = MediaIoBaseDownload(fh, request)
-    done = False
-    while not done:
-        status, done = downloader.next_chunk()
+    while not downloader.next_chunk()[1]:
+        pass
     fh.seek(0)
     with open(local_path, 'wb') as f:
         f.write(fh.read())
 
-# === LIST FILES IN A FOLDER ===
 def list_files(service, folder_id, mime_type='video/mp4'):
     query = f"'{folder_id}' in parents and mimeType='{mime_type}'"
     result = service.files().list(q=query, fields="files(id, name)").execute()
     return result.get('files', [])
 
-# === LOAD REGISTERED FACES ===
 def load_known_faces(service):
     known_faces = []
     reg_files = list_files(service, REG_FOLDER_ID, mime_type='image/jpeg')
@@ -54,24 +50,23 @@ def load_known_faces(service):
         os.remove(fname)
     return known_faces
 
-# === UPLOAD INTRUDER IMAGE ===
 def upload_intruder(service, img_path):
     file_metadata = {'name': os.path.basename(img_path), 'parents': [INTRUDER_FOLDER_ID]}
     media = MediaFileUpload(img_path, mimetype='image/jpeg')
     service.files().create(body=file_metadata, media_body=media, fields='id').execute()
 
-# === BASIC FACE MATCH CHECK ===
-def is_match(face_img, known_faces, threshold=2000):
+def is_match(face_img, known_faces, threshold=MATCH_THRESHOLD):
     face_img = cv2.resize(face_img, (100, 100))
     errors = [np.mean((face_img - kf) ** 2) for kf in known_faces]
+    print(f"[DEBUG] MSE Scores: {errors}")
     return any(err < threshold for err in errors)
 
-# === PROCESS A VIDEO ===
 def process_video(video_path, known_faces, service):
     face_cascade = cv2.CascadeClassifier(CASCADE_PATH)
     video_capture = cv2.VideoCapture(video_path)
-    frame_count = 0
-    detections = 0
+    frame_count, detections = 0, 0
+    uploaded = []
+
     while True:
         ret, frame = video_capture.read()
         if not ret:
@@ -86,35 +81,37 @@ def process_video(video_path, known_faces, service):
                     fname = f"intruder_{ts}.jpg"
                     cv2.imwrite(fname, frame)
                     upload_intruder(service, fname)
+                    uploaded.append(fname)
                     os.remove(fname)
                     detections += 1
         frame_count += 1
     video_capture.release()
-    return detections
+    return detections, uploaded
 
-# === MAIN RUN FUNCTION ===
 def run_detection():
     try:
         drive = init_gdrive()
         known_faces = load_known_faces(drive)
-        total_detections = 0
+        total_detections, intruder_imgs = 0, []
+
         for folder_id in CCTVFEED_IDS:
             videos = list_files(drive, folder_id)
             for v in videos:
                 video_path = f"{v['id']}.mp4"
                 download_file(drive, v['id'], video_path)
-                detections = process_video(video_path, known_faces, drive)
+                detections, uploads = process_video(video_path, known_faces, drive)
                 total_detections += detections
+                intruder_imgs.extend(uploads)
                 os.remove(video_path)
-        return f"✅ Detection complete. {total_detections} intruder(s) uploaded."
+
+        return f"✅ {total_detections} intruder(s) detected.", [f for f in intruder_imgs]
     except Exception as e:
-        return f"❌ Error: {str(e)}"
+        return f"❌ Error: {str(e)}", []
 
-# === GRADIO INTERFACE ===
 gr.Interface(
     fn=run_detection,
     inputs=[],
-    outputs="text",
-    title="CCTV Face Detection",
-    description="Click the button to run detection on uploaded CCTV feeds."
+    outputs=["text", "gallery"],
+    title="CCTV Intruder Detection",
+    description="Click to run detection. Will display and upload unknown faces."
 ).launch()