Test using file paths

app.py

@@ -3,20 +3,29 @@ import gradio as gr
 import numpy as np
 import random
 import PIL
+import face_detection
 
 # import spaces #[uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline
 import torch
 
-# Function to display the uploaded image
-def process_image(image : PIL.Image.Image):
+# # Function to display the uploaded image
+# def process_image(image : PIL.Image.Image):
+#     outputs = face_detection.getCroppedImages(image)
+#     # do AI stuff here
+#     return gr.Image(outputs[0])
+
+def process_image_str(image : str):
+    face_detection.createCroppedSetFromImage(image, "outputs", "imgs")
     # do AI stuff here
-    return image
+    return gr.Image(image)
+    
+    
 
 # Create the Gradio interface
 interface = gr.Interface(
-    fn=process_image, # Function to process the image
-    inputs=gr.Image(type='pil'), # Upload input
+    fn=process_image_str, # Function to process the image
+    inputs=gr.Image(type='filepath'), # Upload input
     outputs=gr.Image(),  # Display output
     allow_flagging='never',
     title="Celebrity Face Detector",

face_detection.py

@@ -0,0 +1,69 @@
+from typing import List
+import face_recognition
+from PIL import Image
+from os import path as p
+import os
+import shutil
+import numpy as np
+
+def getCroppedImages(image: Image.Image, cap = -1):
+    """Takes a PIL image, and returns a list of PIL images with the cropped faces"""
+    image = image.convert("RGB")
+    face_locations = face_recognition.face_locations(np.array(image))
+    
+    outputs = []
+    num = 0
+    for f in face_locations:
+        # allow for capping the number of faces detected, helps prevent multiple faces being detected
+        # when there should only be one
+        if(cap != -1 and num >= cap):
+            break
+        
+        outputs.append(image.crop([f[3], f[0], f[1], f[2]]))
+        num += 1
+        
+    return outputs
+    
+
+def createCroppedSetFromImage(input_path: str, output_dir: str, output_name: str, cap = -1):
+    """
+    input_path: path to the input image
+    output_dir: folder to save the output image to
+    output_name: name for the new image (do not include extension)
+    cap: set to -1 to process all faces detected, otherwise will limit faces to value
+    """
+    imgs = getCroppedImages(Image.open(input_path))
+    num = 0
+    for i in imgs:
+        path = p.join(output_dir, f"{output_name}_{num}.png")
+        i.save(path)
+        print(path)
+        num += 1
+    
+
+def createCroppedSets(input_path: str, output_path: str):
+    """
+    Iterates through subdirectories in the input directory, making a duplicate of it in the output directory.
+    Then iterates through all images in the subdirectory, cropping each face and saving it to the output directory. 
+    
+    For example for params "input" and "output", input/Dwayne Johnson/img.png will be cropped and
+    saved to output/Dwayne Johnson/0_0.png (repeated for all folder and the images they contain) 
+    
+    WARNING: IF A DIRECTORY WITH THE SAME NAME AS THE OUTPUT DIRECTORY, IT WILL BE DELETED
+    """
+    # delete previous folder so there's no conflicts
+    if(p.exists(output_path)):
+        shutil.rmtree(output_path)
+    os.mkdir(output_path)
+    
+    # iterate through all subdirectories in input directory
+    for dir in [name for name in os.listdir(input_path) if p.isdir(p.join(input_path, name))]:
+        sub_out = p.join(output_path, dir)
+        sub_in = p.join(input_path, dir)
+        os.mkdir(sub_out)
+        
+        n = 0
+        # iterate through all files in subdirectory
+        for img in [name for name in os.listdir(sub_in) if p.isfile(p.join(sub_in, name))]:
+            createCroppedSetFromImage(p.join(sub_in, img), sub_out, n, 1)
+            n += 1

requirements.txt

@@ -3,4 +3,6 @@ diffusers
 invisible_watermark
 torch
 transformers
-xformers
+xformers
+face_recognition
+pillow