Update app.py

app.py

@@ -5,138 +5,134 @@ import requests
 import tempfile
 import mediapipe as mp
 import os
-import hashlib
-import io
-from PIL import Image
+import pickle
 from googleapiclient.discovery import build
 from googleapiclient.http import MediaFileUpload
-from google_auth_oauthlib.flow import InstalledAppFlow
 from google.auth.transport.requests import Request
-import pickle
 
 # === CONFIGURATION ===
 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']
+MSE_THRESHOLD = 1000
 
 # === 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 FOR UPLOAD ===
-def get_drive_service_from_token():
-    creds = None
-    if os.path.exists('token.pickle'):
-        with open('token.pickle', 'rb') as token:
-            creds = pickle.load(token)
-    if creds and creds.valid:
-        return build('drive', 'v3', credentials=creds)
-    return None
+def mse(img1, img2):
+    return np.mean((img1.astype("float") - img2.astype("float")) ** 2)
 
-def upload_to_drive_using_token(file_path, filename, folder_id):
-    service = get_drive_service_from_token()
-    if not service:
-        return False
-    file_metadata = {
-        'name': filename,
-        'parents': [folder_id]
-    }
-    media = MediaFileUpload(file_path, resumable=True)
-    service.files().create(body=file_metadata, media_body=media, fields='id').execute()
-    return True
+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)
 
-# === UTILITIES ===
-def get_drive_links(folder_id, mime_prefix="video/"):
+def get_faces_from_drive(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).json()
-    files = res.get("files", [])
-    return [(f["name"], f["id"]) for f in files if f["mimeType"].startswith(mime_prefix)]
-
-def get_file_download_url(file_id):
-    return f"https://drive.google.com/uc?id={file_id}&export=download"
+    results = requests.get(api_url).json().get("files", [])
+    faces = []
+    for f in results:
+        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
+
+def get_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/")]
 
-def download_file(url):
-    response = requests.get(url, stream=True)
+def download_video(link):
+    resp = requests.get(link, stream=True)
     temp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
-    for chunk in response.iter_content(chunk_size=1024 * 1024):
+    for chunk in resp.iter_content(1024*1024):
         temp.write(chunk)
     temp.close()
     return temp.name
 
-def list_registered_faces():
-    st.sidebar.markdown("### 🧑‍💼 Registered Faces")
-    images = get_drive_links(REG_FOLDER_ID, mime_prefix="image/")
-    for name, file_id in images:
-        img_url = get_file_download_url(file_id)
-        img = Image.open(requests.get(img_url, stream=True).raw)
-        st.sidebar.image(img, caption=name, width=120)
-
-# === FACE DETECTION ===
-def detect_faces_from_video(video_path, uploaded_hashes):
-    cap = cv2.VideoCapture(video_path)
-    detected = 0
-    displayed_faces = []
-
-    st.sidebar.markdown("### 🎥 Faces Detected from Video")
-
+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 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
-
-        frame = cv2.resize(frame, (640, 480))
         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
+            for d in results.detections:
+                r = d.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)
-                x, y = max(x, 0), max(y, 0)
-                face = frame[y:y + h, x:x + w]
-                gray = cv2.cvtColor(face, cv2.COLOR_BGR2GRAY)
-                resized = cv2.resize(gray, (100, 100))
-                hsh = hashlib.md5(resized.tobytes()).hexdigest()
-
-                if hsh not in uploaded_hashes:
-                    uploaded_hashes.add(hsh)
-                    tmp_path = os.path.join(tempfile.gettempdir(), f"intruder_{hsh}.jpg")
-                    cv2.imwrite(tmp_path, face)
-                    success = upload_to_drive_using_token(tmp_path, f"intruder_{hsh}.jpg", INTRUDER_FOLDER_ID)
-                    if success:
-                        st.sidebar.image(face, caption="🚨 Intruder", width=120)
-                    else:
-                        st.warning("❌ Upload failed.")
-                else:
-                    st.sidebar.image(face, caption="⛔ Duplicate", width=120)
-                detected += 1
+                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 detected
+    return uploaded
 
 # === STREAMLIT UI ===
-st.set_page_config(page_title="CCTV Face Detector", layout="wide")
-st.title("🔐 Intruder Detection from Google Drive CCTV Feed")
-
-list_registered_faces()
-uploaded_hashes = set()
-
-with st.spinner("Processing videos..."):
-    total_detected = 0
-    for folder_id in CCTVFEED_IDS:
-        videos = get_drive_links(folder_id)
-        for name, vid_id in videos:
-            st.subheader(f"🎞️ {name}")
-            vpath = download_file(get_file_download_url(vid_id))
-            count = detect_faces_from_video(vpath, uploaded_hashes)
-            os.remove(vpath)
-            if count > 0:
-                st.success(f"✅ {count} faces detected in {name}")
-                total_detected += count
-            else:
-                st.warning(f"⚠️ No faces found in {name}")
-
-st.info(f"🎯 Completed processing all videos.")
+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)
+
+intruder_gallery = st.sidebar.container()
+intruder_gallery.header("🚨 Intruders")
+
+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)
+
+st.info(f"🎯 Done. Total intruders uploaded: {total_uploaded}")