Update app.py

app.py

@@ -6,6 +6,7 @@ import tempfile
 import mediapipe as mp
 import os
 import pickle
+import io
 from googleapiclient.discovery import build
 from googleapiclient.http import MediaFileUpload
 from google_auth_oauthlib.flow import InstalledAppFlow
@@ -17,43 +18,49 @@ API_KEY = "AIzaSyDojJrpauA0XZtCCDUuo9xeQHZQamYKsC4"
 CCTVFEED_IDS = ['1KJRkSf2SKEZ1mXS9_si65IwMBtjs6p4n']
 REG_FOLDER_ID = '1qkcR7nQTEtiMH9OFUv2bGxVn08E3dKjF'
 INTRUDER_FOLDER_ID = '1PPAUWU-wMx7fek73p-hqPqYQypYtG8Ob'
-SCOPES = ['https://www.googleapis.com/auth/drive.file']
+FRAME_SKIP = st.sidebar.slider("⏩ Frame Skip", 1, 15, 5)
 
 # === 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 authenticate_drive():
-    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 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)
+def mse(img1, img2):
+    try:
+        if img1.shape != img2.shape:
+            img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))
+        err = np.mean((img1.astype("float") - img2.astype("float")) ** 2)
+        return err
+    except Exception:
+        return float("inf")
+
+def is_duplicate(img, img_list, threshold=200):
+    for existing in img_list:
+        if mse(img, existing) < threshold:
+            return True
+    return False
+
+def get_drive_file_links(folder_id, filter_video=False):
+    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", [])
-    return [(f["name"], f"https://drive.google.com/uc?id={f['id']}&export=download") for f in files if f["mimeType"] in types]
+    links = []
+    for f in files:
+        mime = f["mimeType"]
+        if (not filter_video and mime.startswith("image/")) or (filter_video and mime.startswith("video/")):
+            link = f"https://drive.google.com/uc?id={f['id']}&export=download"
+            links.append((f["name"], link))
+    return links
 
 def download_file(link, suffix):
     response = requests.get(link, 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(image):
+def detect_faces(image):
     rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     results = face_detector.process(rgb)
     faces = []
@@ -61,74 +68,89 @@ def extract_faces(image):
         for det in results.detections:
             bbox = det.location_data.relative_bounding_box
             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)
+            x, y, bw, bh = int(bbox.xmin * w), int(bbox.ymin * h), int(bbox.width * w), int(bbox.height * h)
+            crop = image[max(0, y):y+bh, max(0, x):x+bw]
+            if crop.size > 0:
+                faces.append(crop)
     return faces
 
-def mse(img1, img2):
-    return np.mean((img1.astype("float") - img2.astype("float")) ** 2)
-
-def is_duplicate(new_face, existing_faces):
-    return any(mse(new_face, old) < 200 for old in existing_faces)
+def upload_to_drive(image_np, filename):
+    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)
+    image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
+    buf = io.BytesIO()
+    image_pil.save(buf, format="JPEG")
+    buf.seek(0)
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
+        tmp.write(buf.read())
+        tmp.flush()
+        file_metadata = {'name': filename, 'parents': [INTRUDER_FOLDER_ID]}
+        media = MediaFileUpload(tmp.name, mimetype='image/jpeg')
+        service.files().create(body=file_metadata, media_body=media, fields='id').execute()
+        os.remove(tmp.name)
 
-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)
+def detect_faces_from_video(video_path, registered_faces):
+    cap = cv2.VideoCapture(video_path)
+    detected = []
+    unique = []
+    frame_id = 0
 
-# === MAIN UI ===
-st.set_page_config(layout="wide")
-st.title("🔐 Intruder Detection from CCTV Feed")
-st.sidebar.header("📸 Face Viewer")
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame_id += 1
+        if frame_id % FRAME_SKIP != 0:
+            continue
+        faces = detect_faces(frame)
+        for f in faces:
+            if not any(mse(f, reg) < 200 for reg in registered_faces) and not is_duplicate(f, unique):
+                unique.append(f)
+                upload_to_drive(f, f"intruder_{frame_id}.jpg")
+        detected.extend(faces)
+    cap.release()
+    return detected, unique
 
-drive_service = authenticate_drive()
+# === STREAMLIT UI ===
+st.title("🔐 Intruder Detection")
+st.write("Streams from Google Drive and detects intruders using MediaPipe.")
 
-# === Load Registered Faces ===
+st.sidebar.header("📁 Registered Faces")
+reg_faces_raw = get_drive_file_links(REG_FOLDER_ID, filter_video=False)
 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 = []
-total_detections = 0
+for name, link in reg_faces_raw:
+    img_path = download_file(link, ".jpg")
+    img = cv2.imread(img_path)
+    faces = detect_faces(img)
+    for f in faces:
+        registered_faces.append(f)
+        st.sidebar.image(cv2.cvtColor(f, cv2.COLOR_BGR2RGB), caption=name, width=100)
+    os.remove(img_path)
 
-with st.spinner("📹 Processing videos..."):
+st.sidebar.header("📸 Detected Intruders")
+
+with st.spinner("Processing videos..."):
+    total_uploaded = 0
     for folder_id in CCTVFEED_IDS:
-        videos = get_drive_files(folder_id, ["video/mp4", "video/x-msvideo", "video/avi"])
-        for vname, vlink in videos:
+        vids = get_drive_file_links(folder_id, filter_video=True)
+        for vname, vlink in vids:
             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
-                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):
-                        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(vpath)
-
-st.info(f"🎯 Completed processing. Total intruders uploaded: {total_detections}")
+            vfile = download_file(vlink, ".mp4")
+            detected_faces, unique_faces = detect_faces_from_video(vfile, registered_faces)
+            for i, f in enumerate(unique_faces):
+                st.sidebar.image(cv2.cvtColor(f, cv2.COLOR_BGR2RGB), caption=f"Intruder {i+1}", width=100)
+            if unique_faces:
+                st.success(f"✅ Uploaded {len(unique_faces)} intruder(s).")
+            else:
+                st.info("😎 No new intruders detected.")
+            os.remove(vfile)
+            total_uploaded += len(unique_faces)
+
+st.info(f"🎯 Completed processing all videos.")