Update app.py

app.py

@@ -11,12 +11,12 @@ from PIL import Image
 from facenet_pytorch import MTCNN, InceptionResnetV1
 import json
 import gradio as gr
+from concurrent.futures import ThreadPoolExecutor
 
-# Load Google Drive API credentials from the uploaded JSON file
-SERVICE_ACCOUNT_FILE = './.env'  # Assuming the uploaded JSON file is named '.env'
+# Load Google Drive API credentials
+SERVICE_ACCOUNT_FILE = './.env'  # JSON key file path
 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)
 
@@ -28,29 +28,13 @@ aadhar_folder_id = '1Qtb5DYzSFE67Mbb5ZgDIqWUtdJaDD2F4'
 cphotos_folder_id = '1DGeRqRbCPcfLDdEgP0h5fyX-MF8EQ8AH'
 suspects_folder_id = '1N3RMhVD0OygeufLPYod6IYLtqzvlm3Jv'
 
-# Initialize MTCNN for face detection and InceptionResnetV1 for face recognition
-mtcnn = MTCNN(image_size=160, margin=0, min_face_size=20)
-resnet = InceptionResnetV1(pretrained='vggface2').eval()
-
-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 []
+# Initialize MTCNN and InceptionResnetV1
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+mtcnn = MTCNN(image_size=160, margin=0, min_face_size=20, device=device)
+resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)
 
+# Multithreaded file download
 def download_file(file_id, file_name):
-    print(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()
@@ -61,45 +45,28 @@ def download_file(file_id, file_name):
         file_io.seek(0)
         with open(file_name, 'wb') as f:
             f.write(file_io.read())
-        print(f"Successfully downloaded file: {file_name}")
         return file_name
     except Exception as e:
         print(f"Error downloading file ID: {file_id} - {e}")
         return None
 
-def verify_and_fix_image(image_path):
-    try:
-        with Image.open(image_path) as img:
-            img.verify()
-        with Image.open(image_path) as img:
-            img.save(image_path)
-        print(f"Image verified and cleaned: {image_path}")
-        return True
-    except Exception as e:
-        print(f"Image verification failed for {image_path}. Error: {e}")
-        return False
-
-def upload_file(file_path, folder_id):
-    print(f"Uploading file: {file_path} to folder ID: {folder_id}")
-    file_metadata = {'name': os.path.basename(file_path), 'parents': [folder_id]}
-    media = MediaFileUpload(file_path, resumable=True)
+def list_files_in_folder(folder_id):
+    query = f"'{folder_id}' in parents and trashed=false"
     try:
-        file = drive_service.files().create(body=file_metadata, media_body=media, fields='id').execute()
-        print(f"Uploaded file ID: {file.get('id')}")
-        return file.get('id')
+        results = drive_service.files().list(q=query, pageSize=1000, fields="files(id, name)").execute()
+        return results.get('files', [])
     except Exception as e:
-        print(f"Error uploading file {file_path} to folder ID: {folder_id} - {e}")
-        return None
+        print(f"Error listing files in folder ID: {folder_id} - {e}")
+        return []
 
 def extract_face_embedding(image_path):
     try:
         image = Image.open(image_path)
         face = mtcnn(image)
         if face is not None:
-            embedding = resnet(face.unsqueeze(0))
+            embedding = resnet(face.unsqueeze(0).to(device)).detach().cpu()
             return embedding
         else:
-            print(f"No face detected in image: {image_path}")
             return None
     except Exception as e:
         print(f"Error extracting face embedding for {image_path}. Error: {e}")
@@ -109,23 +76,24 @@ def calculate_similarity(embedding1, embedding2):
     return torch.nn.functional.cosine_similarity(embedding1, embedding2).item()
 
 def process_images(aadhar_folder_id, cphotos_folder_id, suspects_folder_id):
-    print("Starting processing of images...")
     aadhar_files = list_files_in_folder(aadhar_folder_id)
     cphotos_files = list_files_in_folder(cphotos_folder_id)
 
     aadhar_embeddings = []
-    for file in aadhar_files:
-        print(f"Processing Aadhar file: {file['name']}")
-        file_path = download_file(file['id'], file['name'])
-        if file_path and verify_and_fix_image(file_path):
-            embedding = extract_face_embedding(file_path)
-            if embedding is not None:
-                aadhar_embeddings.append((file['name'], embedding))
-
+    
+    # Download and process Aadhar images in parallel
+    with ThreadPoolExecutor() as executor:
+        for file in aadhar_files:
+            file_path = executor.submit(download_file, file['id'], file['name']).result()
+            if file_path:
+                embedding = extract_face_embedding(file_path)
+                if embedding is not None:
+                    aadhar_embeddings.append((file['name'], embedding))
+    
+    # Process CCTV images and find matches in parallel
     for file in cphotos_files:
-        print(f"Processing CCTV file: {file['name']}")
         file_path = download_file(file['id'], file['name'])
-        if not file_path or not verify_and_fix_image(file_path):
+        if not file_path:
             continue
 
         embedding_cctv = extract_face_embedding(file_path)
@@ -137,15 +105,13 @@ def process_images(aadhar_folder_id, cphotos_folder_id, suspects_folder_id):
             similarity = calculate_similarity(embedding_cctv, aadhar_embedding)
             if similarity > 0.8:
                 matched = True
-                print(f"Match found for file: {file['name']} with Aadhar file: {aadhar_name}, Similarity: {similarity:.2f}")
+                print(f"Match found: {file['name']} with Aadhar file: {aadhar_name}, Similarity: {similarity:.2f}")
                 break
 
         if not matched:
             upload_file(file_path, suspects_folder_id)
             print(f"Unmatched image uploaded: {file['name']}")
 
-    print("Processing of images completed.")
-
 def run_script():
     start_time = time.time()
     print("Script started...")