Create app.py

app.py

@@ -0,0 +1,85 @@
+import os
+import pickle
+import numpy as np
+import imgaug.augmenters as iaa
+from PIL import Image
+import face_recognition
+import gradio as gr
+
+def augment_image(image, num_augmented=5):
+    """
+    Apply data augmentation to an image.
+    
+    Parameters:
+    image (PIL.Image): The image to augment.
+    num_augmented (int): Number of augmented images to generate.
+    
+    Returns:
+    List[PIL.Image]: List of augmented images.
+    """
+    image = np.array(image)
+
+    # Define a sequence of augmentation techniques
+    aug = iaa.Sequential([
+        iaa.Fliplr(0.5),  # horizontal flips
+        iaa.Affine(rotate=(-25, 25)),  # rotation
+        iaa.AdditiveGaussianNoise(scale=(0, 0.05*255)),  # noise
+        iaa.Multiply((0.8, 1.2)),  # brightness
+        iaa.GaussianBlur(sigma=(0.0, 1.0))  # blur
+    ])
+
+    # Generate augmented images
+    augmented_images = [Image.fromarray(aug(image=image)) for _ in range(num_augmented)]
+    return augmented_images
+
+def preprocess_and_save_augmented_encodings(images, num_augmented=5):
+    known_encodings = []
+    known_labels = []
+
+    for image in images:
+        # Convert uploaded file to PIL Image
+        original_image = Image.open(image).convert("RGB")  # Ensure the image is in RGB format
+
+        # Augment the image
+        augmented_images = augment_image(original_image, num_augmented=num_augmented)
+
+        # Include the original image in the list of images to encode
+        images_to_encode = [original_image] + augmented_images
+
+        for img in images_to_encode:
+            img_array = np.array(img)
+            # Encode the face
+            encodings = face_recognition.face_encodings(img_array)
+            if encodings:  # Check if an encoding was found
+                encoding = encodings[0]
+                # Store the encoding and the corresponding label
+                known_encodings.append(encoding)
+                known_labels.append(image.name)  # Use the image file name as the label
+
+    # Save encodings and labels to a file
+    output_file = "face_encoding.pkl"
+    data = {"encodings": known_encodings, "labels": known_labels}
+    with open(output_file, "wb") as file:
+        pickle.dump(data, file)
+    
+    return output_file
+
+# Gradio interface
+def gradio_interface(images, num_augmented):
+    pkl_file = preprocess_and_save_augmented_encodings(images, num_augmented=num_augmented)
+    return pkl_file
+
+# Set up Gradio interface
+iface = gr.Interface(
+    fn=gradio_interface,
+    inputs=[
+        gr.inputs.File(type="file", label="Upload Images", multiple=True),
+        gr.inputs.Slider(minimum=1, maximum=10, default=5, label="Number of Augmented Images")
+    ],
+    outputs=gr.File(label="Download Encodings (.pkl)"),
+    title="Face Encoding with Augmentation",
+    description="Upload images, specify the number of augmentations, and download the resulting face encodings in a .pkl file."
+)
+
+# Launch the interface
+iface.launch()