Merge pull request #2 from mochiron-desu/dev-main

Dev main

face_detection.py

@@ -51,25 +51,40 @@ def process_images(image_folder, unique_faces_folder, collection):
                         'face_id': face_id,
                         'image_filename': face_path,
                         'face_encoding': encoded_face,
-                        'occurrences': [filename]
+                        'occurrences': [{
+                            'filename': filename,
+                            'bounding_box': face_location
+                        }]
                     }
 
                     save_face_to_db(collection, face_data)
                     logging.info(f"Saved unique face ID {face_id} to database with path: {face_path}")
 
-                    face_occurrences[face_id] = [filename]
+                    face_occurrences[face_id] = [{
+                        'filename': filename,
+                        'bounding_box': face_location
+                    }]
                     unique_face_count += 1
                 else:
                     logging.info(f"Found existing face ID {existing_face_id} in image: {filename}")
                     # Update the occurrences for the existing face
                     collection.update_one(
                         {'face_id': existing_face_id},
-                        {'$addToSet': {'occurrences': filename}}
+                        {'$push': {'occurrences': {
+                            'filename': filename,
+                            'bounding_box': face_location
+                        }}}
                     )
                     if existing_face_id in face_occurrences:
-                        face_occurrences[existing_face_id].append(filename)
+                        face_occurrences[existing_face_id].append({
+                            'filename': filename,
+                            'bounding_box': face_location
+                        })
                     else:
-                        face_occurrences[existing_face_id] = [filename]
+                        face_occurrences[existing_face_id] = [{
+                            'filename': filename,
+                            'bounding_box': face_location
+                        }]
 
     return unique_face_count, face_occurrences
 

gui.py

@@ -161,11 +161,13 @@ def show_search_results(self, face_id, result):
         logging.info(f"Number of occurrences found: {len(occurrences)}")
 
         def process_images():
-            for i, filename in enumerate(occurrences):
+            for i, occurrence in enumerate(occurrences):
+                filename = occurrence['filename']
+                bounding_box = occurrence['bounding_box']
                 image_path = os.path.join('images', filename)
                 if os.path.exists(image_path):
                     logging.info(f"Processing image: {image_path}")
-                    processed_image = self.process_image_with_face(image_path, face_id)
+                    processed_image = self.process_image_with_face(image_path, bounding_box)
                     self.image_queue.put((processed_image, filename))
                     logging.info(f"Processed image: {filename} ({i + 1}/{len(occurrences)})")
                 else:
@@ -192,23 +194,27 @@ def update_ui():
         result_window.after(100, update_ui)
 
 
-    def process_image_with_face(self, image_path, face_id):
-        # This method now returns the original image instead of processing it
+    def process_image_with_face(self, image_path, bounding_box):
         try:
-            image = face_recognition.load_image_file(image_path)
+            image = Image.open(image_path)
+            draw = ImageDraw.Draw(image)
+
+            # Face_recognition provides (top, right, bottom, left), adjust it to (x0, y0, x1, y1)
+            top, right, bottom, left = bounding_box
 
-            # Convert the image to a PIL image to return it
-            pil_image = Image.fromarray(image)
+            # Draw bounding box
+            draw.rectangle([left, top, right, bottom], outline="yellow", width=30)
 
             # Optionally, you can resize the image if needed
-            pil_image.thumbnail((400, 400))
+            image.thumbnail((400, 400))
 
-            return pil_image
+            return image
         except Exception as e:
             logging.error(f"Error processing image {image_path}: {e}")
             return None
 
 
+
     def display_processed_image(self, parent_frame, pil_image, filename):
         if pil_image:
             frame = tk.Frame(parent_frame, bg="#ffffff", bd=2, relief=tk.GROOVE)