Code optimization

app.py

@@ -7,6 +7,7 @@ import gradio as gr
 from deepface import DeepFace
 from ultralytics import YOLO
 import urllib.request
+import asyncio
 
 with open('config.json', 'r') as f:
     config = json.load(f)
@@ -22,9 +23,9 @@ if not os.path.exists(yolo_weights_filename):
 
 model = YOLO(yolo_weights_filename)
 
-def find_distance(base_face, check_face):
-  result = DeepFace.verify(base_face, check_face, enforce_detection=False)
-  return result['distance']
+async def find_distance(base_face, check_face):
+    result = await asyncio.to_thread(DeepFace.verify, base_face, check_face, enforce_detection=False)
+    return result['distance']
 
 def find_faces(image):
   outputs = model(image)
@@ -39,12 +40,14 @@ def find_faces(image):
       faces.append(crop_img)
   return faces
   
-def load_images_from_zip(zip_path):
+async def load_images_from_zip(zip_path):
     images = []
+    loop = asyncio.get_running_loop()
+    
     with zipfile.ZipFile(zip_path, 'r') as zip_file:
         for file_name in zip_file.namelist():
             with zip_file.open(file_name) as file:
-                img_bytes = file.read()
+                img_bytes = await loop.run_in_executor(None, file.read)
                 img = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR)
                 if img is not None:
                     images.append(img)
@@ -96,26 +99,45 @@ def create_image(images):
 
     return table_cropped
 
-def check(avatars_zip, photos_zip):
-    avatars = [cv2.cvtColor(img, cv2.COLOR_RGB2BGR) for img in load_images_from_zip(avatars_zip.name)]
-    photos = [cv2.cvtColor(img, cv2.COLOR_RGB2BGR) for img in load_images_from_zip(photos_zip.name)]
-
-    input_avatars_faces = [find_faces(avatar) for avatar in avatars]
-    input_avatars_faces = [face for faces in input_avatars_faces for face in faces]
-
+async def process_photo_async(photo, input_avatars_faces):
+    not_found_faces = []
     avatars_faces_count = len(input_avatars_faces)
-
+    input_faces = find_faces(photo)
+    for input_face in input_faces:
+        for i in range(avatars_faces_count):
+            distance = await find_distance(input_avatars_faces[i], input_face)
+            if distance <= FACE_DIST_TRESH:
+                break
+            elif i + 1 == avatars_faces_count:
+                not_found_faces.append(input_face)
+    return not_found_faces
+
+async def check_async(photos, input_avatars_faces, progress):
+    tasks = []
     not_found_faces = []
 
     for photo in photos:
-        input_faces = find_faces(photo)
-        for input_face in input_faces:
-            for i in range(avatars_faces_count):
-                distance = find_distance(input_avatars_faces[i], input_face)
-                if distance <= FACE_DIST_TRESH:
-                    break
-                elif i + 1 == avatars_faces_count:
-                    not_found_faces.append(input_face)
+        task = asyncio.create_task(process_photo_async(photo, input_avatars_faces))
+        tasks.append(task)
+    
+    for i, task in enumerate(tasks):
+        result = await task
+        not_found_faces += result
+        progress((i+1)/len(tasks))
+
+    return not_found_faces
+
+def check(avatars_zip, photos_zip, progress=gr.Progress()):
+    avatars = asyncio.run(load_images_from_zip(avatars_zip.name))
+    avatars = [cv2.cvtColor(avatar, cv2.COLOR_RGB2BGR) for avatar in avatars]
+
+    photos = asyncio.run(load_images_from_zip(photos_zip.name))
+    photos = [cv2.cvtColor(photo, cv2.COLOR_RGB2BGR) for photo in photos]
+
+    input_avatars_faces = [find_faces(avatar) for avatar in avatars]
+    input_avatars_faces = [face for faces in input_avatars_faces for face in faces]
+
+    not_found_faces = asyncio.run(check_async(photos, input_avatars_faces, progress))
 
     return create_image(not_found_faces)
 
@@ -134,4 +156,4 @@ with gr.Blocks(theme='soft', title='SquadDetective') as blocks:
     outputs=gr.outputs.Image(type='numpy', label='Report')
     process_button.click(fn=check, inputs=inputs, outputs=outputs)
 
-blocks.launch()
+blocks.queue(concurrency_count=1).launch()