Update app.py

app.py

@@ -1,3 +1,4 @@
+
 import streamlit as st
 import cv2
 import numpy as np
@@ -6,133 +7,134 @@ import tempfile
 import mediapipe as mp
 import os
 import pickle
+import io
 from googleapiclient.discovery import build
-from googleapiclient.http import MediaFileUpload
+from googleapiclient.http import MediaIoBaseUpload
+from google_auth_oauthlib.flow import InstalledAppFlow
 from google.auth.transport.requests import Request
+from PIL import Image
 
 # === CONFIGURATION ===
 API_KEY = "AIzaSyDojJrpauA0XZtCCDUuo9xeQHZQamYKsC4"
 CCTVFEED_IDS = ['1KJRkSf2SKEZ1mXS9_si65IwMBtjs6p4n']
 REG_FOLDER_ID = '1qkcR7nQTEtiMH9OFUv2bGxVn08E3dKjF'
 INTRUDER_FOLDER_ID = '1PPAUWU-wMx7fek73p-hqPqYQypYtG8Ob'
-MSE_THRESHOLD = 1000
+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)
 
-def mse(img1, img2):
-    return np.mean((img1.astype("float") - img2.astype("float")) ** 2)
+# === AUTH ===
+def get_drive_service_with_pickle():
+    creds = None
+    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)
+            creds = flow.run_local_server(port=0)
+        with open("token.pickle", "wb") as token:
+            pickle.dump(creds, token)
+    return build("drive", "v3", credentials=creds)
 
-def download_image(url):
-    resp = requests.get(url, stream=True)
-    arr = np.asarray(bytearray(resp.content), dtype=np.uint8)
-    return cv2.imdecode(arr, cv2.IMREAD_COLOR)
+drive_service = get_drive_service_with_pickle()
 
-def get_faces_from_drive(folder_id):
+# === 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)"
-    results = requests.get(api_url).json().get("files", [])
-    faces = []
-    for f in results:
+    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"
-            image = download_image(url)
-            rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            detections = face_detector.process(rgb)
-            if detections.detections:
-                for d in detections.detections:
-                    r = d.location_data.relative_bounding_box
-                    ih, iw, _ = image.shape
-                    x, y, w, h = int(r.xmin * iw), int(r.ymin * ih), int(r.width * iw), int(r.height * ih)
-                    face = image[y:y+h, x:x+w]
-                    if face.size:
-                        face = cv2.resize(face, (100, 100))
-                        faces.append(face)
-    return faces
+            img_resp = requests.get(url)
+            img = Image.open(io.BytesIO(img_resp.content)).convert("RGB")
+            images.append(np.array(img))
+    return images
 
-def get_video_links(folder_id):
+def get_drive_video_links(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)"
-    results = requests.get(api_url).json().get("files", [])
-    return [(f["name"], f"https://drive.google.com/uc?id={f['id']}&export=download") 
-            for f in results if f["mimeType"].startswith("video/")]
+    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/")]
 
 def download_video(link):
-    resp = requests.get(link, stream=True)
+    response = requests.get(link, stream=True)
     temp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
-    for chunk in resp.iter_content(1024*1024):
+    for chunk in response.iter_content(chunk_size=1024 * 1024):
         temp.write(chunk)
     temp.close()
     return temp.name
 
-def upload_intruder(img):
-    creds = None
-    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())
-    service = build('drive', 'v3', credentials=creds)
-    temp = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg")
-    cv2.imwrite(temp.name, img)
-    file_metadata = {'name': os.path.basename(temp.name), 'parents': [INTRUDER_FOLDER_ID]}
-    media = MediaFileUpload(temp.name, mimetype='image/jpeg')
-    service.files().create(body=file_metadata, media_body=media).execute()
-    os.remove(temp.name)
+def faces_from_frame(frame):
+    rgb = cv2.cvtColor(frame, 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)))
+    return faces
+
+def mse(img1, img2):
+    return np.mean((img1.astype("float") - img2.astype("float")) ** 2)
 
-def process_video(path, reg_faces, intruder_gallery):
-    cap = cv2.VideoCapture(path)
-    uploaded_faces = []
-    uploaded = 0
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        results = face_detector.process(rgb)
-        if results.detections:
-            for d in results.detections:
-                r = d.location_data.relative_bounding_box
-                ih, iw, _ = frame.shape
-                x, y, w, h = int(r.xmin * iw), int(r.ymin * ih), int(r.width * iw), int(r.height * ih)
-                face = frame[y:y+h, x:x+w]
-                if face.size:
-                    face_resized = cv2.resize(face, (100, 100))
-                    if not any(mse(face_resized, r) < MSE_THRESHOLD for r in reg_faces + uploaded_faces):
-                        uploaded_faces.append(face_resized)
-                        intruder_gallery.image(face, caption="👀 New Intruder", use_column_width=True)
-                        try:
-                            upload_intruder(face)
-                            uploaded += 1
-                        except:
-                            st.warning("Upload failed.")
-    cap.release()
-    return uploaded
+def is_duplicate(face, uploaded_faces, threshold=200):
+    return any(mse(face, f) < threshold for f in uploaded_faces)
 
-# === STREAMLIT UI ===
-st.set_page_config(layout="wide")
-st.title("🔐 Intruder Detection (MediaPipe + Google Drive)")
-sidebar = st.sidebar
-sidebar.header("🧍 Registered Faces")
-registered_faces = get_faces_from_drive(REG_FOLDER_ID)
-for idx, face in enumerate(registered_faces):
-    sidebar.image(cv2.cvtColor(face, cv2.COLOR_BGR2RGB), caption=f"Reg {idx+1}", use_column_width=True)
+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
 
-intruder_gallery = st.sidebar.container()
-intruder_gallery.header("🚨 Intruders")
+# === MAIN APP ===
+st.set_page_config(page_title="CCTV Face Intruder Detection", layout="wide")
+st.title("🔐 CCTV Intruder Face Detection (MediaPipe + Google Drive)")
 
-total_uploaded = 0
-for folder_id in CCTVFEED_IDS:
-    videos = get_video_links(folder_id)
-    for vname, link in videos:
-        st.subheader(f"🎥 {vname}")
-        with st.spinner("Downloading video..."):
-            video_file = download_video(link)
-        with st.spinner("Detecting faces..."):
-            uploaded = process_video(video_file, registered_faces, intruder_gallery)
-            st.success(f"✅ Uploaded {uploaded} intruder(s) from {vname}")
-            total_uploaded += uploaded
-        os.remove(video_file)
+registered_faces = get_drive_images(REG_FOLDER_ID)
+st.sidebar.header("📌 Registered Faces")
+for img in registered_faces:
+    st.sidebar.image(img, width=100)
 
-st.info(f"🎯 Done. Total intruders uploaded: {total_uploaded}")
+uploaded_intruders = []
+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)
+            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")
+                    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)
+            cap.release()
+            os.remove(path)
+    st.success("✅ Done.")