Update app.py

app.py

@@ -2,6 +2,9 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 import os
 import time
 import json
+from google.oauth2 import service_account
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaIoBaseDownload, MediaFileUpload
 
 # Load Google Drive API credentials from the uploaded JSON file
 SERVICE_ACCOUNT_FILE = './.env'  # Assuming the uploaded JSON file is named '.env'
@@ -19,25 +22,13 @@ aadhar_folder_id = '1Qtb5DYzSFE67Mbb5ZgDIqWUtdJaDD2F4'
 cphotos_folder_id = '1DGeRqRbCPcfLDdEgP0h5fyX-MF8EQ8AH'
 suspects_folder_id = '1N3RMhVD0OygeufLPYod6IYLtqzvlm3Jv'
 
-
+# Increase the efficiency of ThreadPoolExecutor by balancing the number of threads
+MAX_WORKERS = 10
 
 # Improved function to batch process images with better face comparison logic
-def list_files_in_folder(folder_id):
-    print(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, owners, parents)").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.")
-        else:
-            print(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}")
-        return []
-
 def batch_process_images(aadhar_folder_id, cphotos_folder_id, suspects_folder_id):
+    start_time = time.time()
+
     # Fetch Aadhar and CCTV files
     aadhar_files = list_files_in_folder(aadhar_folder_id)
     cphotos_files = list_files_in_folder(cphotos_folder_id)
@@ -45,17 +36,17 @@ def batch_process_images(aadhar_folder_id, cphotos_folder_id, suspects_folder_id
     # Extract Aadhar embeddings in parallel
     aadhar_embeddings = []
     print("Extracting embeddings from Aadhar images...")
-    with ThreadPoolExecutor() as executor:
-        future_to_aadhar = {executor.submit(download_file, file['id'], file['name']): file for file in aadhar_files}
+
+    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
+        future_to_aadhar = {executor.submit(download_and_extract_embedding, file): file for file in aadhar_files}
         for future in as_completed(future_to_aadhar):
-            file = future_to_aadhar[future]
             try:
-                file_path = future.result()
-                if file_path:
-                    embedding = extract_face_embedding(file_path)
-                    if embedding is not None:
-                        aadhar_embeddings.append((file['name'], embedding))
+                result = future.result()
+                if result:
+                    file_name, embedding = result
+                    aadhar_embeddings.append((file_name, embedding))
             except Exception as e:
+                file = future_to_aadhar[future]
                 print(f"Error processing Aadhar file {file['name']}: {e}")
 
     # Verify that embeddings were successfully extracted for at least one Aadhar file
@@ -66,30 +57,29 @@ def batch_process_images(aadhar_folder_id, cphotos_folder_id, suspects_folder_id
     # Process CCTV files and compare with Aadhar embeddings
     unmatched_files = []
     print("Processing CCTV images and comparing embeddings...")
-    with ThreadPoolExecutor() as executor:
-        future_to_cphoto = {executor.submit(download_file, file['id'], file['name']): file for file in cphotos_files}
+
+    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
+        future_to_cphoto = {executor.submit(download_and_extract_embedding, file): file for file in cphotos_files}
         for future in as_completed(future_to_cphoto):
-            file = future_to_cphoto[future]
             try:
-                file_path = future.result()
-                if not file_path:
+                result = future.result()
+                if not result:
                     continue
 
-                embedding_cctv = extract_face_embedding(file_path)
-                if embedding_cctv is None:
-                    continue
+                file_name, embedding_cctv = result
 
                 # Compare CCTV embedding with all Aadhaar embeddings
                 matched = compare_faces_with_embeddings(embedding_cctv, aadhar_embeddings)
                 if not matched:
-                    unmatched_files.append(file_path)
-                    print(f"Unmatched image queued for upload: {file['name']}")
+                    unmatched_files.append(file_name)
+                    print(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}")
 
     # Upload unmatched files to suspect folder
     print("Uploading unmatched images to the suspect folder...")
-    with ThreadPoolExecutor() as executor:
+    with ThreadPoolExecutor(max_workers=MAX_WORKERS) 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:
@@ -98,6 +88,23 @@ def batch_process_images(aadhar_folder_id, cphotos_folder_id, suspects_folder_id
             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(file):
+    """
+    Download the file and extract its face embedding.
+    This function combines the download and extraction step to reduce overhead.
+    """
+    try:
+        file_path = download_file(file['id'], file['name'])
+        if file_path:
+            embedding = extract_face_embedding(file_path)
+            return file['name'], embedding
+    except Exception as e:
+        print(f"Error processing file {file['name']}: {e}")
+    return None
+
 def compare_faces_with_embeddings(embedding_cctv, aadhar_embeddings, threshold=0.75):
     """
     Compare the given CCTV embedding with a list of Aadhar embeddings.