Update app.py

app.py

@@ -5,123 +5,120 @@ import requests
 import tempfile
 import mediapipe as mp
 import os
+from PIL import Image
+from io import BytesIO
 
 # === CONFIGURATION ===
 API_KEY = "AIzaSyDojJrpauA0XZtCCDUuo9xeQHZQamYKsC4"
 CCTVFEED_IDS = ['1KJRkSf2SKEZ1mXS9_si65IwMBtjs6p4n']
 REG_FOLDER_ID = '1qkcR7nQTEtiMH9OFUv2bGxVn08E3dKjF'
 INTRUDER_FOLDER_ID = '1PPAUWU-wMx7fek73p-hqPqYQypYtG8Ob'
-FACE_SIZE = (100, 100)
-MATCH_THRESHOLD = 6000
 
 # === 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)
 
-def get_drive_links(folder_id, mime_filter):
-    api_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(api_url)
-    files = response.json().get("files", [])
-    return [(f["name"], f"https://drive.google.com/uc?id={f['id']}&export=download") 
-            for f in files if f["mimeType"] in mime_filter]
+# === FUNCTIONS ===
+def list_drive_files(folder_id):
+    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)
+    return response.json().get("files", [])
 
-def download_file(link, suffix):
-    response = requests.get(link, stream=True)
+def stream_drive_file(file_id, suffix):
+    url = f"https://drive.google.com/uc?id={file_id}&export=download"
+    response = requests.get(url, stream=True)
     temp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
-    for chunk in response.iter_content(chunk_size=1024 * 1024):
+    for chunk in response.iter_content(1024 * 1024):
         temp.write(chunk)
     temp.close()
     return temp.name
 
-def extract_faces_from_image(image_path):
-    img = cv2.imread(image_path)
-    faces = []
-    if img is None: return []
-    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    results = face_detector.process(rgb)
-    if results.detections:
-        for det in results.detections:
-            bboxC = det.location_data.relative_bounding_box
-            ih, iw, _ = img.shape
-            x, y, w, h = int(bboxC.xmin * iw), int(bboxC.ymin * ih), \
-                         int(bboxC.width * iw), int(bboxC.height * ih)
-            face = img[y:y + h, x:x + w]
-            if face.size > 0:
-                face = cv2.resize(cv2.cvtColor(face, cv2.COLOR_BGR2GRAY), FACE_SIZE)
-                faces.append(face)
-    return faces
+def mse(img1, img2):
+    if img1.shape != img2.shape:
+        return np.inf
+    return np.mean((img1 - img2) ** 2)
 
 def load_registered_faces():
     reg_faces = []
-    images = get_drive_links(REG_FOLDER_ID, ['image/jpeg', 'image/png'])
-    st.sidebar.write(f"📁 Registered Images: {len(images)}")
-    for name, link in images:
-        img_path = download_file(link, ".jpg")
-        faces = extract_faces_from_image(img_path)
-        reg_faces.extend(faces)
-        os.remove(img_path)
+    files = list_drive_files(REG_FOLDER_ID)
+    for f in files:
+        if "image" in f['mimeType']:
+            path = stream_drive_file(f['id'], ".jpg")
+            img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+            img = cv2.resize(img, (100, 100))
+            reg_faces.append(img)
+            os.remove(path)
     return reg_faces
 
-def mse(imageA, imageB):
-    return np.mean((imageA.astype("float") - imageB.astype("float")) ** 2)
+def upload_intruder_image(img_np, intruder_name="intruder_face.jpg"):
+    # Convert image to JPEG in memory
+    img = Image.fromarray(cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB))
+    buffer = BytesIO()
+    img.save(buffer, format="JPEG")
+    buffer.seek(0)
+    
+    # Upload to intruder folder using resumable upload
+    upload_url = f"https://www.googleapis.com/upload/drive/v3/files?uploadType=multipart&key={API_KEY}"
+    metadata = {
+        "name": intruder_name,
+        "parents": [INTRUDER_FOLDER_ID]
+    }
+    files = {
+        "metadata": ("metadata.json", BytesIO(str.encode(str(metadata).replace("'", '"'))), "application/json"),
+        "file": ("image.jpg", buffer, "image/jpeg")
+    }
+    response = requests.post(upload_url, files=files)
+    return response.ok
 
-def is_intruder(face, known_faces):
-    errors = [mse(face, kf) for kf in known_faces]
-    return all(err > MATCH_THRESHOLD for err in errors)
-
-def extract_and_match_faces_from_video(video_path, known_faces):
+def detect_faces_from_video(video_path, reg_faces):
     cap = cv2.VideoCapture(video_path)
-    detected_faces = []
-    intruders = []
+    detected, intruders, sidebar_faces = 0, 0, []
+
     while cap.isOpened():
         ret, frame = cap.read()
-        if not ret: break
+        if not ret:
+            break
         rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         results = face_detector.process(rgb)
         if results.detections:
             for det in results.detections:
                 bbox = det.location_data.relative_bounding_box
-                ih, iw, _ = frame.shape
-                x, y, w, h = int(bbox.xmin * iw), int(bbox.ymin * ih), \
-                             int(bbox.width * iw), int(bbox.height * ih)
-                face = frame[y:y + h, x:x + w]
-                if face.size > 0:
-                    gray_face = cv2.resize(cv2.cvtColor(face, cv2.COLOR_BGR2GRAY), FACE_SIZE)
-                    detected_faces.append(face)
-                    if is_intruder(gray_face, known_faces):
-                        intruders.append(face)
-        if len(detected_faces) >= 5:
-            break
+                h, w, _ = frame.shape
+                x, y, bw, bh = int(bbox.xmin * w), int(bbox.ymin * h), int(bbox.width * w), int(bbox.height * h)
+                face = frame[y:y+bh, x:x+bw]
+                if face.size == 0:
+                    continue
+                face_gray = cv2.cvtColor(cv2.resize(face, (100, 100)), cv2.COLOR_BGR2GRAY)
+                sidebar_faces.append(face)
+                is_known = any(mse(face_gray, kf) < 2000 for kf in reg_faces)
+                if not is_known:
+                    upload_intruder_image(face, f"intruder_{intruders}.jpg")
+                    intruders += 1
+                    detected += 1
     cap.release()
-    return detected_faces, intruders
-
-def show_faces(title, face_list):
-    st.sidebar.subheader(title)
-    for face in face_list:
-        st.sidebar.image(cv2.cvtColor(face, cv2.COLOR_BGR2RGB), width=100)
+    return detected, sidebar_faces
 
-# === STREAMLIT APP ===
-st.title("🔐 Intruder Detection (with Face Match)")
-st.write("Detect and compare faces from CCTV videos streamed from Google Drive.")
+# === STREAMLIT UI ===
+st.set_page_config(layout="wide")
+st.title("🔐 Intruder Detection from CCTV Feed")
+st.write("Streams public videos from Google Drive, compares detected faces with registered images, and uploads unknown faces to the intruder folder.")
 
-with st.spinner("🚀 Loading registered faces..."):
-    known_faces = load_registered_faces()
+reg_faces = load_registered_faces()
+st.sidebar.markdown("### Faces Detected")
 
-with st.spinner("📽 Processing CCTV videos..."):
-    all_detected, all_intruders = [], []
+with st.spinner("Processing CCTV feed..."):
+    total_intruders = 0
     for folder_id in CCTVFEED_IDS:
-        videos = get_drive_links(folder_id, ['video/mp4', 'video/x-msvideo', 'video/avi'])
-        for vname, link in videos:
-            st.subheader(f"🎥 {vname}")
-            video_file = download_file(link, ".mp4")
-            detected, intruders = extract_and_match_faces_from_video(video_file, known_faces)
-            os.remove(video_file)
-            st.success(f"✅ Detected {len(detected)} faces | 🚨 Intruders: {len(intruders)}")
-            all_detected.extend(detected)
-            all_intruders.extend(intruders)
-
-show_faces("👥 Faces from Video", all_detected)
-show_faces("🚨 Intruders", all_intruders)
+        files = list_drive_files(folder_id)
+        for f in files:
+            if f['mimeType'].startswith('video/'):
+                st.subheader(f"🎥 {f['name']}")
+                path = stream_drive_file(f['id'], ".mp4")
+                intruders, detected_faces = detect_faces_from_video(path, reg_faces)
+                total_intruders += intruders
+                for face in detected_faces:
+                    st.sidebar.image(face, width=100)
+                os.remove(path)
 
-st.info(f"🎯 Total videos scanned: {len(CCTVFEED_IDS)} | Total intruders detected: {len(all_intruders)}")
+st.success(f"✅ Intruders found and uploaded: {total_intruders}")