Update app.py

app.py

@@ -5,6 +5,10 @@ import requests
 import tempfile
 import mediapipe as mp
 import os
+import pickle
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaFileUpload
+from PIL import Image
 
 # === CONFIGURATION ===
 API_KEY = "AIzaSyDojJrpauA0XZtCCDUuo9xeQHZQamYKsC4"
@@ -12,133 +16,123 @@ CCTVFEED_IDS = ['1KJRkSf2SKEZ1mXS9_si65IwMBtjs6p4n']
 REG_FOLDER_ID = '1qkcR7nQTEtiMH9OFUv2bGxVn08E3dKjF'
 INTRUDER_FOLDER_ID = '1PPAUWU-wMx7fek73p-hqPqYQypYtG8Ob'
 
-# === SETUP ===
 mp_face_detection = mp.solutions.face_detection
-mp_drawing = mp.solutions.drawing_utils
 face_detector = mp_face_detection.FaceDetection(model_selection=0, min_detection_confidence=0.5)
 
-# === UTILS ===
-def get_drive_files(folder_id, mime_prefix="image/"):
+# === INTRUDER UPLOAD ===
+def get_drive_service():
+    try:
+        with open("token.pickle", "rb") as token:
+            creds = pickle.load(token)
+        return build("drive", "v3", credentials=creds)
+    except Exception as e:
+        print(f"❌ Drive auth failed: {e}")
+        return None
+
+def upload_intruder_face(image_path):
+    service = get_drive_service()
+    if service is None:
+        print("🔒 Upload skipped (auth failed)")
+        return
+
+    metadata = {"name": os.path.basename(image_path), "parents": [INTRUDER_FOLDER_ID]}
+    media = MediaFileUpload(image_path, mimetype="image/jpeg")
+    try:
+        service.files().create(body=metadata, media_body=media, fields="id").execute()
+        print(f"✅ Uploaded to INTRUDER folder: {image_path}")
+    except Exception as e:
+        print(f"❌ Upload failed: {e}")
+
+# === HELPERS ===
+def get_drive_files(folder_id, allowed_mime):
     url = f"https://www.googleapis.com/drive/v3/files?q='{folder_id}'+in+parents&key={API_KEY}&fields=files(id,name,mimeType)"
-    r = requests.get(url)
-    files = r.json().get("files", [])
-    return [f for f in files if f["mimeType"].startswith(mime_prefix)]
+    response = requests.get(url)
+    files = response.json().get("files", [])
+    return [(f["name"], f["id"]) for f in files if f["mimeType"] in allowed_mime]
 
-def get_image_from_drive(file_id):
+def download_file(file_id, suffix):
     url = f"https://drive.google.com/uc?id={file_id}&export=download"
-    r = requests.get(url)
-    img_array = np.asarray(bytearray(r.content), dtype=np.uint8)
-    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
-    return img
+    response = requests.get(url, stream=True)
+    temp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+    for chunk in response.iter_content(chunk_size=1024 * 1024):
+        temp.write(chunk)
+    temp.close()
+    return temp.name
 
-def extract_face_embeddings(img):
-    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+def extract_face_embedding(image):
+    rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     results = face_detector.process(rgb)
-    faces = []
     if results.detections:
         for det in results.detections:
             bbox = det.location_data.relative_bounding_box
-            h, w, _ = img.shape
-            x, y = int(bbox.xmin * w), int(bbox.ymin * h)
-            bw, bh = int(bbox.width * w), int(bbox.height * h)
-            face_crop = img[max(0, y):y+bh, max(0, x):x+bw]
-            if face_crop.size > 0:
-                resized = cv2.resize(face_crop, (128, 128))
-                faces.append(resized)
-    return faces
-
-def compare_faces(face1, face2):
-    try:
-        diff = cv2.absdiff(face1, face2)
-        return np.mean(diff)
-    except:
-        return 9999
-
-def upload_intruder_image(img_array, name_hint="intruder"):
-    _, buffer = cv2.imencode('.jpg', img_array)
-    img_bytes = buffer.tobytes()
-    upload_url = f"https://www.googleapis.com/upload/drive/v3/files?uploadType=multipart&key={API_KEY}"
-
-    metadata = {
-        "name": f"{name_hint}.jpg",
-        "parents": [INTRUDER_FOLDER_ID]
-    }
-
-    boundary = "foo_bar_baz"
-    headers = {
-        "Content-Type": f"multipart/related; boundary={boundary}"
-    }
-
-    meta_json = str(metadata).replace("'", '"')
-    body_prefix = (
-        "--" + boundary + "\r\n"
-        "Content-Type: application/json; charset=UTF-8\r\n\r\n"
-        + meta_json + "\r\n"
-        "--" + boundary + "\r\n"
-        "Content-Type: image/jpeg\r\n\r\n"
-    ).encode("utf-8")
-
-    body_suffix = ("\r\n--" + boundary + "--").encode("utf-8")
-    multipart_body = body_prefix + img_bytes + body_suffix
-
-    r = requests.post(upload_url, headers=headers, data=multipart_body)
-    return r.status_code == 200
-
-def download_video(link):
-    r = requests.get(link, stream=True)
-    temp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
-    for chunk in r.iter_content(chunk_size=1024*1024):
-        temp.write(chunk)
-    temp.close()
-    return temp.name
-
-# === MAIN ===
+            h, w, _ = image.shape
+            x1 = int(bbox.xmin * w)
+            y1 = int(bbox.ymin * h)
+            x2 = int((bbox.xmin + bbox.width) * w)
+            y2 = int((bbox.ymin + bbox.height) * h)
+            face_crop = image[y1:y2, x1:x2]
+            return cv2.resize(face_crop, (128, 128))
+    return None
+
+def is_match(face1, face2):
+    if face1 is None or face2 is None:
+        return False
+    diff = np.linalg.norm(face1.astype("float") - face2.astype("float"))
+    return diff < 30.0
+
+# === STREAMLIT UI ===
+st.set_page_config(page_title="Intruder Detection", layout="wide")
 st.title("🔐 Intruder Detection from CCTV Feed")
-st.write("Streaming videos from Google Drive, detecting faces using MediaPipe, and uploading unmatched intruders.")
-
-# Load registered faces
-st.sidebar.header("✅ Registered Faces")
-registered_faces = []
-reg_files = get_drive_files(REG_FOLDER_ID, "image/")
-for f in reg_files:
-    img = get_image_from_drive(f['id'])
-    faces = extract_face_embeddings(img)
-    registered_faces.extend(faces)
-    for face in faces:
-        st.sidebar.image(cv2.cvtColor(face, cv2.COLOR_BGR2RGB), caption=f['name'], width=100)
-
-# Process CCTV videos
-with st.spinner("Processing CCTV videos..."):
-    total_detections = 0
+st.write("Streaming videos from Google Drive, detecting and matching faces.")
+
+sidebar_faces = []
+
+with st.spinner("Loading registered faces..."):
+    registered = []
+    reg_images = get_drive_files(REG_FOLDER_ID, ["image/jpeg", "image/png"])
+    for name, fid in reg_images:
+        img_path = download_file(fid, ".jpg")
+        img = cv2.imread(img_path)
+        emb = extract_face_embedding(img)
+        if emb is not None:
+            registered.append((name, emb))
+
+with st.spinner("Scanning CCTV feed..."):
+    total_intruders = 0
     for folder_id in CCTVFEED_IDS:
-        url = f"https://www.googleapis.com/drive/v3/files?q='{folder_id}'+in+parents&key={API_KEY}&fields=files(id,name,mimeType)"
-        response = requests.get(url)
-        files = response.json().get("files", [])
-        for f in files:
-            if f["mimeType"].startswith("video/"):
-                st.subheader(f"🎥 {f['name']}")
-                link = f"https://drive.google.com/uc?id={f['id']}&export=download"
-                video_file = download_video(link)
-                cap = cv2.VideoCapture(video_file)
-                found_intruder = False
-
-                while cap.isOpened():
-                    ret, frame = cap.read()
-                    if not ret:
-                        break
-                    faces = extract_face_embeddings(frame)
-                    for face in faces:
-                        st.sidebar.image(cv2.cvtColor(face, cv2.COLOR_BGR2RGB), caption="Detected", width=100)
-                        distances = [compare_faces(face, reg_face) for reg_face in registered_faces]
-                        if not distances or min(distances) > 45:  # Distance threshold
-                            uploaded = upload_intruder_image(face, f"intruder_{f['name']}")
-                            if uploaded:
-                                st.warning("⚠️ Intruder uploaded to Google Drive!")
-                            else:
-                                st.error("Failed to upload intruder.")
-                            found_intruder = True
-                cap.release()
-                os.remove(video_file)
-                if not found_intruder:
-                    st.success("✅ No intruders found.")
-    st.info("🎯 Completed processing all videos.")
+        videos = get_drive_files(folder_id, ["video/mp4", "video/x-msvideo", "video/avi"])
+        for vname, fid in videos:
+            st.subheader(f"🎥 {vname}")
+            vid_path = download_file(fid, ".mp4")
+            cap = cv2.VideoCapture(vid_path)
+            frame_count = 0
+            matched_faces = 0
+            unmatched_faces = 0
+
+            while cap.isOpened():
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                face_img = extract_face_embedding(frame)
+                if face_img is not None:
+                    match_found = any(is_match(face_img, emb) for _, emb in registered)
+                    if match_found:
+                        matched_faces += 1
+                    else:
+                        unmatched_faces += 1
+                        temp = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg")
+                        cv2.imwrite(temp.name, face_img)
+                        upload_intruder_face(temp.name)
+                        sidebar_faces.append(temp.name)
+                frame_count += 1
+            cap.release()
+
+            st.success(f"🎯 Processed {frame_count} frames — Matched: {matched_faces}, Unmatched: {unmatched_faces}")
+            total_intruders += unmatched_faces
+
+st.info(f"🎯 Completed processing all videos. Intruders detected: {total_intruders}")
+
+with st.sidebar:
+    st.markdown("## 🚨 Intruder Snapshots")
+    for face_path in sidebar_faces:
+        st.image(face_path, width=160)