Update app.py

app.py

@@ -1,5 +1,3 @@
-
-
 import os
 import cv2
 from google.oauth2 import service_account
@@ -7,11 +5,14 @@ from googleapiclient.discovery import build
 from googleapiclient.http import MediaIoBaseDownload, MediaFileUpload
 import io
 import time
+import torch
+import numpy as np
 from PIL import Image
+from facenet_pytorch import MTCNN, InceptionResnetV1
 
-# Load Google Drive API credentials
+# Load Google Drive API credentials from the uploaded JSON file
+SERVICE_ACCOUNT_FILE = './salesforce-api-439514-d6b432a2e20e.json'  # Assuming the JSON file is in the root directory
 SCOPES = ['https://www.googleapis.com/auth/drive']
-SERVICE_ACCOUNT_FILE = r'C:\Users\Ajaya\Downloads\salesforce-api-439514-d6b432a2e20e.json'  # Use raw string to avoid Unicode escape issues
 credentials = service_account.Credentials.from_service_account_file(SERVICE_ACCOUNT_FILE, scopes=SCOPES)
 drive_service = build('drive', 'v3', credentials=credentials)
 
@@ -20,6 +21,10 @@ aadhar_folder_id = '1jXklVbxTni1gnduuKgJ7zI6zVks6s0LP'
 cphotos_folder_id = '1JVJMtHuIHxZyxAekT6ZQFkEWePOimIlX'
 suspects_folder_id = '1LJE5aw_g5Jvqd5cLZ7UiJ8LyR36JkBKP'
 
+# 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"
@@ -30,8 +35,6 @@ def list_files_in_folder(folder_id):
             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}.")
-            for file in files:
-                print(f"Found file: Name={file['name']}, ID={file['id']}, Type={file['mimeType']}, Owner={file['owners'][0]['displayName']}, Parent Folder ID={file['parents']}")
         return files
     except Exception as e:
         print(f"Error listing files in folder ID: {folder_id} - {e}")
@@ -83,62 +86,62 @@ def upload_file(file_path, folder_id):
         print(f"Error uploading file {file_path} to folder ID: {folder_id} - {e}")
         return None
 
-def detect_face(image_path):
-    print(f"Detecting faces in image: {image_path}")
-    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-    image = cv2.imread(image_path)
-    if image is None:
-        print(f"Error reading image: {image_path}. The file may not be a valid image.")
-        return False
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
-    face_detected = len(faces) > 0
-    print(f"Faces detected: {face_detected} in image: {image_path}")
-    return face_detected
+def extract_face_embedding(image_path):
+    try:
+        image = Image.open(image_path)
+        face = mtcnn(image)  # Detect and crop the face
+        if face is not None:
+            embedding = resnet(face.unsqueeze(0))  # Get the face embedding
+            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}")
+        return None
+
+def calculate_similarity(embedding1, embedding2):
+    # Calculate cosine similarity between two embeddings
+    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_faces = []
+    # Extract embeddings for all Aadhar images
+    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) and detect_face(file_path):
-            img = cv2.imread(file_path)
-            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-            face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-            faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
-            if len(faces) > 0:
-                aadhar_faces.append((file['name'], faces, img))
+        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))
 
+    # Compare each CCTV image with all Aadhar embeddings
     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):
             continue
+
+        embedding_cctv = extract_face_embedding(file_path)
+        if embedding_cctv is None:
+            continue
+
         matched = False
-        if detect_face(file_path):
-            img_c = cv2.imread(file_path)
-            gray_c = cv2.cvtColor(img_c, cv2.COLOR_BGR2GRAY)
-            face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-            faces_c = face_cascade.detectMultiScale(gray_c, scaleFactor=1.1, minNeighbors=5)
-            if len(faces_c) > 0:
-                for aadhar_name, aadhar_face_regions, aadhar_img in aadhar_faces:
-                    for (x, y, w, h) in aadhar_face_regions:
-                        aadhar_face = aadhar_img[y:y+h, x:x+w]
-                        for (xc, yc, wc, hc) in faces_c:
-                            cctv_face = img_c[yc:yc+hc, xc:xc+wc]
-                            if aadhar_face.shape == cctv_face.shape and not (cv2.subtract(aadhar_face, cctv_face).any()):
-                                matched = True
-                                print(f"Match found for file: {file['name']} with Aadhar file: {aadhar_name}")
-                                break
-                    if matched:
-                        break
+        for aadhar_name, aadhar_embedding in aadhar_embeddings:
+            similarity = calculate_similarity(embedding_cctv, aadhar_embedding)
+            if similarity > 0.8:  # Adjust the threshold as necessary
+                matched = True
+                print(f"Match found for file: {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 main():