Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import streamlit as st
 import cv2
 import numpy as np
@@ -7,9 +6,8 @@ import tempfile
 import mediapipe as mp
 import os
 import pickle
-import io
 from googleapiclient.discovery import build
-from googleapiclient.http import MediaIoBaseUpload
+from googleapiclient.http import MediaFileUpload
 from google_auth_oauthlib.flow import InstalledAppFlow
 from google.auth.transport.requests import Request
 from PIL import Image
@@ -23,118 +21,114 @@ SCOPES = ['https://www.googleapis.com/auth/drive.file']
 
 # === 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)
 
-# === AUTH ===
-def get_drive_service_with_pickle():
+def authenticate_drive():
     creds = None
-    if os.path.exists("token.pickle"):
-        with open("token.pickle", "rb") as token:
+    if os.path.exists('token.pickle'):
+        with open('token.pickle', 'rb') as token:
             creds = pickle.load(token)
     if not creds or not creds.valid:
         if creds and creds.expired and creds.refresh_token:
             creds.refresh(Request())
         else:
-            flow = InstalledAppFlow.from_client_secrets_file("credentials.json", SCOPES)
+            flow = InstalledAppFlow.from_client_secrets_file('credentials.json', SCOPES)
             creds = flow.run_local_server(port=0)
-        with open("token.pickle", "wb") as token:
+        with open('token.pickle', 'wb') as token:
             pickle.dump(creds, token)
-    return build("drive", "v3", credentials=creds)
-
-drive_service = get_drive_service_with_pickle()
-
-# === UTILITY FUNCTIONS ===
-def get_drive_images(folder_id):
-    api_url = f"https://www.googleapis.com/drive/v3/files?q='{folder_id}'+in+parents&key={API_KEY}&fields=files(id,name,mimeType)"
-    res = requests.get(api_url)
-    files = res.json().get("files", [])
-    images = []
-    for f in files:
-        if f["mimeType"].startswith("image/"):
-            url = f"https://drive.google.com/uc?id={f['id']}&export=download"
-            img_resp = requests.get(url)
-            img = Image.open(io.BytesIO(img_resp.content)).convert("RGB")
-            images.append(np.array(img))
-    return images
+    return build('drive', 'v3', credentials=creds)
 
-def get_drive_video_links(folder_id):
+def get_drive_files(folder_id, types):
     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"].startswith("video/")]
+    return [(f["name"], f"https://drive.google.com/uc?id={f['id']}&export=download") for f in files if f["mimeType"] in types]
 
-def download_video(link):
+def download_file(link, suffix):
     response = requests.get(link, stream=True)
-    temp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
-    for chunk in response.iter_content(chunk_size=1024 * 1024):
+    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 faces_from_frame(frame):
-    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+def extract_faces(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, _ = frame.shape
-            x, y = int(bbox.xmin * w), int(bbox.ymin * h)
-            width, height = int(bbox.width * w), int(bbox.height * h)
-            face = frame[y:y + height, x:x + width]
-            if face.size > 0:
-                faces.append(cv2.resize(face, (128, 128)))
+            h, w, _ = image.shape
+            x, y, w_, h_ = int(bbox.xmin * w), int(bbox.ymin * h), int(bbox.width * w), int(bbox.height * h)
+            x, y = max(x, 0), max(y, 0)
+            cropped = image[y:y+h_, x:x+w_]
+            if cropped.size:
+                face = cv2.resize(cropped, (128, 128))
+                faces.append(face)
     return faces
 
 def mse(img1, img2):
     return np.mean((img1.astype("float") - img2.astype("float")) ** 2)
 
-def is_duplicate(face, uploaded_faces, threshold=200):
-    return any(mse(face, f) < threshold for f in uploaded_faces)
+def is_duplicate(new_face, existing_faces):
+    return any(mse(new_face, old) < 200 for old in existing_faces)
 
-def upload_intruder_image(image_np):
-    try:
-        img = Image.fromarray(image_np)
-        buf = io.BytesIO()
-        img.save(buf, format='JPEG')
-        buf.seek(0)
-        media = MediaIoBaseUpload(buf, mimetype="image/jpeg")
-        drive_service.files().create(
-            body={"name": "intruder.jpg", "parents": [INTRUDER_FOLDER_ID]},
-            media_body=media,
-            fields="id"
-        ).execute()
-        return True
-    except Exception as e:
-        st.error(f"Upload error: {e}")
-        return False
+def upload_image_to_drive(service, folder_id, image_array, name):
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg")
+    cv2.imwrite(temp_file.name, image_array)
+    file_metadata = {'name': name, 'parents': [folder_id]}
+    media = MediaFileUpload(temp_file.name, mimetype='image/jpeg')
+    service.files().create(body=file_metadata, media_body=media, fields='id').execute()
+    os.remove(temp_file.name)
 
-# === MAIN APP ===
-st.set_page_config(page_title="CCTV Face Intruder Detection", layout="wide")
-st.title("🔐 CCTV Intruder Face Detection (MediaPipe + Google Drive)")
+# === MAIN UI ===
+st.set_page_config(layout="wide")
+st.title("🔐 Intruder Detection from CCTV Feed")
+st.sidebar.header("📸 Face Viewer")
 
-registered_faces = get_drive_images(REG_FOLDER_ID)
-st.sidebar.header("📌 Registered Faces")
-for img in registered_faces:
-    st.sidebar.image(img, width=100)
+drive_service = authenticate_drive()
 
+# === Load Registered Faces ===
+registered_faces = []
+with st.spinner("📂 Loading registered faces..."):
+    reg_files = get_drive_files(REG_FOLDER_ID, ["image/jpeg", "image/png"])
+    for name, link in reg_files:
+        img_path = download_file(link, ".jpg")
+        img = cv2.imread(img_path)
+        faces = extract_faces(img)
+        for f in faces:
+            registered_faces.append(f)
+            st.sidebar.image(cv2.cvtColor(f, cv2.COLOR_BGR2RGB), caption=f"REG: {name}", width=100)
+        os.remove(img_path)
+
+# === Process Videos ===
 uploaded_intruders = []
-with st.spinner("🔍 Processing videos..."):
+total_detections = 0
+
+with st.spinner("📹 Processing videos..."):
     for folder_id in CCTVFEED_IDS:
-        for name, link in get_drive_video_links(folder_id):
-            st.subheader(f"🎥 {name}")
-            path = download_video(link)
-            cap = cv2.VideoCapture(path)
+        videos = get_drive_files(folder_id, ["video/mp4", "video/x-msvideo", "video/avi"])
+        for vname, vlink in videos:
+            st.subheader(f"🎥 {vname}")
+            vpath = download_file(vlink, ".mp4")
+            cap = cv2.VideoCapture(vpath)
+            frame_id = 0
             while cap.isOpened():
                 ret, frame = cap.read()
                 if not ret:
                     break
-                detected_faces = faces_from_frame(frame)
-                for f in detected_faces:
-                    st.sidebar.image(f, width=100, caption="🎯 Detected")
+                frame_id += 1
+                faces = extract_faces(frame)
+                for idx, f in enumerate(faces):
+                    st.sidebar.image(cv2.cvtColor(f, cv2.COLOR_BGR2RGB), caption=f"VID Frame {frame_id}", width=100)
                     if not any(mse(f, reg) < 200 for reg in registered_faces) and not is_duplicate(f, uploaded_intruders):
-                        if upload_intruder_image(f):
-                            uploaded_intruders.append(f)
+                        upload_image_to_drive(drive_service, INTRUDER_FOLDER_ID, f, f"intruder_{vname}_{frame_id}_{idx}.jpg")
+                        st.success(f"⚠️ Intruder detected in frame {frame_id}")
+                        uploaded_intruders.append(f)
+                        total_detections += 1
             cap.release()
-            os.remove(path)
-    st.success("✅ Done.")
+            os.remove(vpath)
+
+st.info(f"🎯 Completed processing. Total intruders uploaded: {total_detections}")