Update app.py

app.py

@@ -1,119 +1,94 @@
-import os
 import cv2
-import face_recognition
-import numpy as np
-import time
-import pickle
+import mediapipe as mp
+import os
+import tempfile
+import shutil
+import requests
+from bs4 import BeautifulSoup
 import gradio as gr
-from googleapiclient.discovery import build
-from googleapiclient.http import MediaIoBaseDownload, MediaFileUpload
-from io import BytesIO
-
-# === CONFIGURATION ===
-CCTVFEED_IDS = ['1KJRkSf2SKEZ1mXS9_si65IwMBtjs6p4n']
-REG_FOLDER_ID = '1qkcR7nQTEtiMH9OFUv2bGxVn08E3dKjF'
-INTRUDER_FOLDER_ID = '1PPAUWU-wMx7fek73p-hqPqYQypYtG8Ob'
-TOKEN_PICKLE = 'token.pickle'
 
-def init_gdrive():
-    with open(TOKEN_PICKLE, 'rb') as token:
-        creds = pickle.load(token)
-    return build('drive', 'v3', credentials=creds)
+# MediaPipe setup
+mp_face_detection = mp.solutions.face_detection
+mp_drawing = mp.solutions.drawing_utils
 
-def download_file(service, file_id, local_path):
-    request = service.files().get_media(fileId=file_id)
-    fh = BytesIO()
-    downloader = MediaIoBaseDownload(fh, request)
-    while not downloader.next_chunk()[1]:
-        pass
-    fh.seek(0)
-    with open(local_path, 'wb') as f:
-        f.write(fh.read())
+def download_public_gdrive_file(file_id, save_dir):
+    try:
+        url = f"https://drive.google.com/uc?id={file_id}&export=download"
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            ext = ".avi" if "avi" in response.headers.get('Content-Disposition', '') else ".mp4"
+            output_path = os.path.join(save_dir, f"{file_id}{ext}")
+            with open(output_path, 'wb') as f:
+                f.write(response.content)
+            return output_path
+    except Exception as e:
+        print(f"Error downloading {file_id}: {e}")
+    return None
 
-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', [])
+def get_video_ids_from_folder(folder_url):
+    response = requests.get(folder_url)
+    soup = BeautifulSoup(response.text, "html.parser")
+    video_ids = []
+    for link in soup.find_all("a"):
+        href = link.get("href", "")
+        if "/file/d/" in href:
+            video_id = href.split("/file/d/")[1].split("/")[0]
+            video_ids.append(video_id)
+    return list(set(video_ids))
 
-def load_known_faces(service):
-    known_encodings = []
-    files = list_files(service, REG_FOLDER_ID, mime_type='image/jpeg')
-    print(f"[INFO] Found {len(files)} registered images.")
-    for f in files:
-        fname = f"{f['id']}.jpg"
-        download_file(service, f['id'], fname)
-        image = face_recognition.load_image_file(fname)
-        encodings = face_recognition.face_encodings(image)
-        if encodings:
-            known_encodings.append(encodings[0])
-        os.remove(fname)
-    return known_encodings
+def detect_faces_from_video(video_path):
+    face_count = 0
+    temp_dir = tempfile.mkdtemp()
+    frames_with_faces = []
 
-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()
+    cap = cv2.VideoCapture(video_path)
+    frame_index = 0
 
-def is_match(encoding, known_encodings, tolerance=0.45):
-    results = face_recognition.compare_faces(known_encodings, encoding, tolerance)
-    return any(results)
+    with mp_face_detection.FaceDetection(model_selection=0, min_detection_confidence=0.5) as detector:
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frame_index += 1
+            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            results = detector.process(rgb_frame)
 
-def process_video(video_path, known_encodings, service):
-    cap = cv2.VideoCapture(video_path)
-    frame_count, detections = 0, 0
-    uploaded = []
+            if results.detections:
+                face_count += len(results.detections)
+                for detection in results.detections:
+                    mp_drawing.draw_detection(frame, detection)
 
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        if frame_count % 30 == 0:
-            rgb = frame[:, :, ::-1]
-            locations = face_recognition.face_locations(rgb)
-            encodings = face_recognition.face_encodings(rgb, locations)
-            print(f"[FRAME {frame_count}] Found {len(encodings)} face(s).")
-            for encoding in encodings:
-                if not is_match(encoding, known_encodings):
-                    ts = int(time.time())
-                    fname = f"intruder_{ts}_{frame_count}.jpg"
-                    cv2.imwrite(fname, frame)
-                    upload_intruder(service, fname)
-                    uploaded.append(fname)
-                    os.remove(fname)
-                    detections += 1
-        frame_count += 1
+                frame_path = os.path.join(temp_dir, f"frame_{frame_index}.jpg")
+                cv2.imwrite(frame_path, frame)
+                frames_with_faces.append(frame_path)
 
     cap.release()
-    print(f"[INFO] Done with {video_path}, {frame_count} frames processed.")
-    return detections, uploaded
+    return frames_with_faces, face_count
 
-def run_detection():
-    try:
-        service = init_gdrive()
-        known_encodings = load_known_faces(service)
-        total_detected, total_videos, intruder_imgs = 0, 0, []
+def run_on_drive_folder(folder_url):
+    folder_url = folder_url.strip()
+    video_ids = get_video_ids_from_folder(folder_url)
+    if not video_ids:
+        return "❌ No videos found or access denied.", []
+    
+    frames_all = []
+    total_faces = 0
+    for vid_id in video_ids[:5]:  # Limit to 5 videos to manage processing time
+        video_path = download_public_gdrive_file(vid_id, "/tmp")
+        if video_path:
+            frames, count = detect_faces_from_video(video_path)
+            frames_all.extend(frames)
+            total_faces += count
 
-        for folder_id in CCTVFEED_IDS:
-            videos = list_files(service, folder_id)
-            total_videos += len(videos)
-            print(f"[INFO] Found {len(videos)} video(s) in {folder_id}.")
-            for v in videos:
-                path = f"{v['id']}.mp4"
-                download_file(service, v['id'], path)
-                detected, found = process_video(path, known_encodings, service)
-                total_detected += detected
-                intruder_imgs.extend(found)
-                os.remove(path)
+    return f"✅ Detected {total_faces} face(s) across {len(frames_all)} frame(s) from {len(video_ids)} video(s).", frames_all
 
-        summary = f"✅ Processed {total_videos} video(s). {total_detected} intruder(s) detected."
-        return summary, intruder_imgs
-    except Exception as e:
-        return f"❌ Error: {str(e)}", []
+demo = gr.Interface(
+    fn=run_on_drive_folder,
+    inputs=gr.Textbox(label="Google Drive Folder URL (Public)"),
+    outputs=[gr.Textbox(label="Summary"), gr.Gallery(label="Detected Frames")],
+    title="Intruder Detection from Google Drive Folder (MP4 + AVI)",
+    description="Paste the link to a *public* Google Drive folder containing .mp4 or .avi videos."
+)
 
-gr.Interface(
-    fn=run_detection,
-    inputs=[],
-    outputs=["text", "gallery"],
-    title="CCTV Intruder Detector (GPU Ready)",
-    description="Detects intruders by comparing faces in CCTV footage with registered staff/guest photos."
-).launch()
+if __name__ == "__main__":
+    demo.launch()