Update app.py

app.py

@@ -1,166 +1,177 @@
-from concurrent.futures import ThreadPoolExecutor, as_completed
 import os
-import time
 import json
+import logging
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from google.oauth2 import service_account
 from googleapiclient.discovery import build
-from googleapiclient.http import MediaIoBaseDownload
-import io
-import numpy as np
+from googleapiclient.http import MediaIoBaseDownload, MediaFileUpload
+from PIL import Image
+import torch
 
-# Load Google Drive API credentials from the uploaded JSON file
+# Set up logging
+logging.basicConfig(level=logging.DEBUG)
+
+# Load Google Drive API credentials
 SERVICE_ACCOUNT_FILE = './.env'  # Assuming the uploaded JSON file is named '.env'
 SCOPES = ['https://www.googleapis.com/auth/drive']
 
-# Load credentials from the JSON file
-with open(SERVICE_ACCOUNT_FILE, 'r') as json_file:
-    credentials_info = json.load(json_file)
-
-credentials = service_account.Credentials.from_service_account_info(credentials_info, scopes=SCOPES)
-drive_service = build('drive', 'v3', credentials=credentials)
+try:
+    if not os.path.exists(SERVICE_ACCOUNT_FILE):
+        logging.error(f"Service account file '{SERVICE_ACCOUNT_FILE}' not found.")
+        raise FileNotFoundError(f"Service account file '{SERVICE_ACCOUNT_FILE}' not found.")
+
+    with open(SERVICE_ACCOUNT_FILE, 'r') as json_file:
+        credentials_info = json.load(json_file)
+    credentials = service_account.Credentials.from_service_account_info(credentials_info, scopes=SCOPES)
+    drive_service = build('drive', 'v3', credentials=credentials)
+    logging.debug("Google Drive API credentials loaded successfully.")
+except json.JSONDecodeError as e:
+    logging.error(f"Error parsing JSON service account file: {e}")
+except FileNotFoundError as e:
+    logging.error(f"Service account file missing: {e}")
+except Exception as e:
+    logging.error(f"Unexpected error while loading credentials: {e}")
+    raise
 
 # Folder IDs
 aadhar_folder_id = '1Qtb5DYzSFE67Mbb5ZgDIqWUtdJaDD2F4'
 cphotos_folder_id = '1DGeRqRbCPcfLDdEgP0h5fyX-MF8EQ8AH'
 suspects_folder_id = '1N3RMhVD0OygeufLPYod6IYLtqzvlm3Jv'
 
-MAX_WORKERS = 10
+# Lazy loaded models (initialized only when needed)
+mtcnn = None
+resnet = None
+
+def initialize_models():
+    global mtcnn, resnet
+    if mtcnn is None or resnet is None:
+        logging.debug("Loading face detection and recognition models...")
+        from facenet_pytorch import MTCNN, InceptionResnetV1
+        mtcnn = MTCNN(image_size=160, margin=0, min_face_size=20)
+        resnet = InceptionResnetV1(pretrained='vggface2').eval()
+        logging.debug("Models loaded successfully.")
 
 def list_files_in_folder(folder_id):
-    print(f"Listing files in folder ID: {folder_id}")
+    logging.debug(f"Listing files in folder ID: {folder_id}")
     query = f"'{folder_id}' in parents and trashed=false"
     try:
         results = drive_service.files().list(q=query, pageSize=1000, fields="files(id, name, mimeType)").execute()
         files = results.get('files', [])
         if not files:
-            print(f"No files found in folder ID: {folder_id}. Ensure that the folder has files and the service account has access.")
+            logging.warning(f"No files found in folder ID: {folder_id}. Ensure that the folder has files and the service account has access.")
         else:
-            print(f"Found {len(files)} files in folder ID: {folder_id}.")
+            logging.debug(f"Found {len(files)} files in folder ID: {folder_id}.")
         return files
     except Exception as e:
-        print(f"Error listing files in folder ID: {folder_id} - {e}")
+        logging.error(f"Error listing files in folder ID: {folder_id} - {e}")
         return []
 
-# Improved function to batch process images with better face comparison logic
-def batch_process_images(aadhar_folder_id, cphotos_folder_id, suspects_folder_id):
-    start_time = time.time()
+def download_file(file_id, file_name):
+    logging.debug(f"Attempting to download file: {file_name} (ID: {file_id})")
+    if not file_name.lower().endswith(('.jpg', '.jpeg', '.png')):
+        file_name += ".jpg"
+    try:
+        request = drive_service.files().get_media(fileId=file_id)
+        file_io = io.BytesIO()
+        downloader = MediaIoBaseDownload(file_io, request)
+        done = False
+        while not done:
+            _, done = downloader.next_chunk()
+        file_io.seek(0)
+        with open(file_name, 'wb') as f:
+            f.write(file_io.read())
+        logging.debug(f"Successfully downloaded file: {file_name}")
+        return file_name
+    except Exception as e:
+        logging.error(f"Error downloading file ID: {file_id} - {e}")
+        return None
 
-    # Fetch Aadhar and CCTV files
+def extract_face_embedding(image_path):
+    try:
+        initialize_models()  # Load models if not already loaded
+        image = Image.open(image_path)
+        face = mtcnn(image)
+        if face is not None:
+            embedding = resnet(face.unsqueeze(0))
+            return embedding
+        else:
+            logging.warning(f"No face detected in image: {image_path}")
+            return None
+    except Exception as e:
+        logging.error(f"Error extracting face embedding for {image_path}. Error: {e}")
+        return None
+
+def calculate_similarity(embedding1, embedding2):
+    return torch.nn.functional.cosine_similarity(embedding1, embedding2).item()
+
+def compare_faces_with_embeddings(embedding_cctv, aadhar_embeddings, threshold=0.7):
+    for aadhar_name, aadhar_embedding in aadhar_embeddings:
+        similarity = calculate_similarity(embedding_cctv, aadhar_embedding)
+        if similarity > threshold:
+            logging.info(f"Match found: Similarity = {similarity:.2f} with Aadhar file: {aadhar_name}")
+            return True
+    return False
+
+def batch_process_images(aadhar_folder_id, cphotos_folder_id, suspects_folder_id):
+    logging.debug("Starting processing of images...")
     aadhar_files = list_files_in_folder(aadhar_folder_id)
     cphotos_files = list_files_in_folder(cphotos_folder_id)
 
-    # Extract Aadhar embeddings in parallel
     aadhar_embeddings = []
-    print("Extracting embeddings from Aadhar images...")
-
-    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
-        future_to_aadhar = {executor.submit(download_and_extract_embedding_from_memory, file): file for file in aadhar_files}
+    with ThreadPoolExecutor() as executor:
+        future_to_aadhar = {executor.submit(download_file, file['id'], file['name']): file for file in aadhar_files}
         for future in as_completed(future_to_aadhar):
+            file = future_to_aadhar[future]
             try:
-                result = future.result()
-                if result:
-                    file_name, embedding = result
-                    aadhar_embeddings.append(embedding)
+                file_path = future.result()
+                if file_path:
+                    embedding = extract_face_embedding(file_path)
+                    if embedding is not None:
+                        aadhar_embeddings.append((file['name'], embedding))
             except Exception as e:
-                file = future_to_aadhar[future]
-                print(f"Error processing Aadhar file {file['name']}: {e}")
+                logging.error(f"Error processing Aadhar file {file['name']}: {e}")
 
-    # Verify that embeddings were successfully extracted for at least one Aadhar file
-    if len(aadhar_embeddings) == 0:
-        print("No valid Aadhar face embeddings found. Exiting.")
+    if not aadhar_embeddings:
+        logging.warning("No valid Aadhar face embeddings found. Exiting.")
         return
 
-    # Stack all embeddings for efficient vectorized comparison
-    aadhar_embeddings = np.vstack(aadhar_embeddings)
-
-    # Process CCTV files and compare with Aadhar embeddings
     unmatched_files = []
-    print("Processing CCTV images and comparing embeddings...")
-
-    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
-        future_to_cphoto = {executor.submit(download_and_extract_embedding_from_memory, file): file for file in cphotos_files}
+    with ThreadPoolExecutor() as executor:
+        future_to_cphoto = {executor.submit(download_file, file['id'], file['name']): file for file in cphotos_files}
         for future in as_completed(future_to_cphoto):
+            file = future_to_cphoto[future]
             try:
-                result = future.result()
-                if not result:
+                file_path = future.result()
+                if not file_path:
                     continue
 
-                file_name, embedding_cctv = result
+                embedding_cctv = extract_face_embedding(file_path)
+                if embedding_cctv is None:
+                    continue
 
-                # Compare CCTV embedding with all Aadhaar embeddings
-                matched = vectorized_compare_faces_with_embeddings(embedding_cctv, aadhar_embeddings)
+                matched = compare_faces_with_embeddings(embedding_cctv, aadhar_embeddings)
                 if not matched:
-                    unmatched_files.append(file_name)
-                    print(f"Unmatched image queued for upload: {file_name}")
+                    unmatched_files.append(file_path)
+                    logging.info(f"Unmatched image queued for upload: {file['name']}")
             except Exception as e:
-                file = future_to_cphoto[future]
-                print(f"Error processing CCTV file {file['name']}: {e}")
+                logging.error(f"Error processing CCTV file {file['name']}: {e}")
 
-    # Upload unmatched files to suspect folder
-    print("Uploading unmatched images to the suspect folder...")
-    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
-        upload_futures = [executor.submit(upload_file, file_name, suspects_folder_id) for file_name in unmatched_files]
+    with ThreadPoolExecutor() as executor:
+        upload_futures = [executor.submit(upload_file, file_path, suspects_folder_id) for file_path in unmatched_files]
         for future in as_completed(upload_futures):
             try:
                 future.result()
-                print(f"Successfully uploaded unmatched image to suspects folder.")
+                logging.info(f"Successfully uploaded unmatched image to suspects folder.")
             except Exception as e:
-                print(f"Error uploading unmatched file: {e}")
-
-    end_time = time.time()
-    print(f"Total processing time: {end_time - start_time:.2f} seconds.")
-
-def download_and_extract_embedding_from_memory(file):
-    """
-    Download the file in memory and extract its face embedding.
-    """
-    file_stream = download_file_in_memory(file['id'])
-    if file_stream:
-        embedding = extract_face_embedding_from_stream(file_stream)
-        if embedding is not None:
-            return file['name'], embedding
-    return None
-
-def download_file_in_memory(file_id):
-    try:
-        request = drive_service.files().get_media(fileId=file_id)
-        file_io = io.BytesIO()
-        downloader = MediaIoBaseDownload(file_io, request)
-        done = False
-        while not done:
-            _, done = downloader.next_chunk()
-        file_io.seek(0)
-        return file_io
-    except Exception as e:
-        print(f"Error downloading file ID: {file_id} - {e}")
-        return None
-
-def extract_face_embedding_from_stream(file_stream):
-    try:
-        image = Image.open(file_stream)
-        face = mtcnn(image)  # Detect and crop the face
-        if face is not None:
-            embedding = resnet(face.unsqueeze(0))  # Get the face embedding
-            return embedding.detach().numpy()
-    except Exception as e:
-        print(f"Error extracting face embedding from stream: {e}")
-    return None
-
-def vectorized_compare_faces_with_embeddings(embedding_cctv, aadhar_embeddings, threshold=0.75):
-    """
-    Compare the given CCTV embedding with a list of Aadhar embeddings in a vectorized manner.
-    """
-    similarities = np.dot(aadhar_embeddings, embedding_cctv.T) / (np.linalg.norm(aadhar_embeddings, axis=1) * np.linalg.norm(embedding_cctv))
-    max_similarity = np.max(similarities)
-    if max_similarity > threshold:
-        print(f"Match found: Similarity = {max_similarity:.2f}")
-        return True
-    return False
+                logging.error(f"Error uploading unmatched file: {e}")
 
 def main():
-    print("Script started...")
+    start_time = time.time()
+    logging.debug("Script started...")
     batch_process_images(aadhar_folder_id, cphotos_folder_id, suspects_folder_id)
+    end_time = time.time()
+    logging.debug(f"Script completed in {end_time - start_time:.2f} seconds.")
 
 if __name__ == "__main__":
-    main()
-
+    main()