Update app.py

app.py

@@ -196,54 +196,6 @@ def add_hats_sticker(img_bgr, sticker_path, faces):
 
     return img_with_stickers
 
-def add_headbands_sticker(img_bgr, sticker_path, faces):
-    headband_pil = Image.open(sticker_path)
-
-    # Check the color mode and convert to RGBA
-    headband_rgba = headband_pil.convert('RGBA')
-    
-    # Convert the headband_rgba to BGRA
-    r, g, b, a = headband_rgba.split()
-    headband_bgra = Image.merge("RGBA", (b, g, r, a))
-    
-    # A copy of the original image
-    img_with_stickers = img_bgr.copy()
-
-    for face in faces:
-        landmarks = face_landmarking(img_bgr, face)
-
-        # Determine the forehead region using landmarks
-        # Assuming the headband will be placed between the temples
-        left_temple = landmarks.part(0)   
-        right_temple = landmarks.part(16) 
-
-        # Calculate the width of the headband based on the temples
-        headband_width = int(abs(left_temple.x - right_temple.x) * 1.6)
-        headband_height = int(headband_width * headband_bgra.height / headband_bgra.width)
-
-        # Resize the headband image
-        resized_headband = headband_bgra.resize((headband_width, headband_height))
-
-        # Calculate the angle of tilt using the eyes as reference
-        left_eye_indices = range(36, 42)
-        right_eye_indices = range(42, 48)
-        angle = calculate_eye_angle(landmarks, left_eye_indices, right_eye_indices)
-
-        # Rotate the headband image
-        rotated_headband = resized_headband.rotate(-angle, expand=True, resample=Image.BICUBIC)
-
-        # Calculate the position for the headband
-        x1 = (left_temple.x + right_temple.x) // 2 - (headband_width // 2)
-        y1 = left_temple.y - (headband_height // 2) 
-
-        # Convert PIL image to NumPy array
-        headband_np = np.array(rotated_headband)
-
-        # Overlay the headband on the image
-        img_with_stickers = overlay_transparent(img_with_stickers, headband_np, x1, y1)
-
-    return img_with_stickers
-
 
 # Function to add glasses stickers
 def add_glasses_sticker(img_bgr, sticker_path, faces):
@@ -410,83 +362,11 @@ def add_animal_faces_sticker(img_bgr, sticker_path, faces):
     return img_with_stickers
 
 
-# Function to process the image
-def process_image(image, sticker_choice):
-    if sticker_choice:
-        # add .png to the sticker_choice
-        sticker_name = sticker_choice + '.png'
-        # find sticker's category
-        sticker_category = STICKER_TO_CATEGORY[sticker_name]
-        # Path to the single sticker
-        sticker_path = os.path.join('stickers',sticker_category, sticker_name)
-
-        # Convert PIL image to OpenCV format BGR
-        image_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-
-        # Detect faces
-        faces = face_detecting(image_bgr)
-        # print(sticker_category)
-        if sticker_category == 'ears':
-            img_with_stickers = add_ears_sticker(image_bgr, sticker_path, faces)
-        elif sticker_category == 'glasses':
-            img_with_stickers = add_glasses_sticker(image_bgr, sticker_path, faces)
-        elif sticker_category == 'noses':
-            img_with_stickers = add_noses_sticker(image_bgr, sticker_path, faces)
-        elif sticker_category == 'headbands':
-            img_with_stickers = add_ears_sticker(image_bgr, sticker_path, faces)
-        else:
-            img_with_stickers = add_glasses_sticker(image_bgr, sticker_path, faces)
-        # Convert back to PIL image
-        img_with_stickers_pil = Image.fromarray(cv2.cvtColor(img_with_stickers, cv2.COLOR_BGR2RGB))
-        return img_with_stickers_pil
-    else:
-        return image
-
-def process_image_with_selections(image_input):
-    # Convert PIL image to OpenCV format BGR
-    image_bgr = cv2.cvtColor(np.array(image_input), cv2.COLOR_RGB2BGR)
-
-    # Detect faces
-    faces = face_detecting(image_bgr)
-    
-    # A copy of the original image to apply stickers on
-    img_with_stickers = image_bgr.copy()
-
-    for category, sticker_name in sticker_selections.items():
-        if sticker_name:  # Check if a sticker was selected in this category
-            # the sticker file path
-            sticker_path = os.path.join('stickers', category, sticker_name + '.png')
-
-            # Apply the selected sticker based on its category
-            if category == 'ears':
-                img_with_stickers = add_ears_sticker(img_with_stickers, sticker_path, faces)
-            elif category == 'glasses':
-                img_with_stickers = add_glasses_sticker(img_with_stickers, sticker_path, faces)
-            elif category == 'noses':
-                img_with_stickers = add_noses_sticker(img_with_stickers, sticker_path, faces)
-            elif category == 'headbands':
-                img_with_stickers = add_hats_sticker(img_with_stickers, sticker_path, faces)
-            elif category == 'hats':
-                img_with_stickers = add_hats_sticker(img_with_stickers, sticker_path, faces)
-            elif category == 'animal face':
-                img_with_stickers = add_animal_faces_sticker(img_with_stickers, sticker_path, faces)
-            else:
-                img_with_stickers = img_with_stickers
-    # Convert back to PIL image
-    img_with_stickers_pil = Image.fromarray(cv2.cvtColor(img_with_stickers, cv2.COLOR_BGR2RGB))
-    
-    print("Selected stickers:")
-    for category, selection in sticker_selections.items():
-        print(f"{category}: {selection}")
-
-    return img_with_stickers_pil
-
 # This dictionary will hold the user's sticker selections
-sticker_selections = {}
+# sticker_selections = {}
 
 # Function to update sticker selections
 def update_selections(category, selection):
-    # sticker_selections[category] = selection
     sticker_selections[category] = None if selection == "None" else selection
     return ""
 
@@ -494,7 +374,6 @@ def update_selections(category, selection):
 def load_example_image(image_path):
     return gr.Image.from_file(image_path)
 
-from PIL import Image
 
 def resize_image(image, target_width, target_height):
     # Maintain aspect ratio
@@ -518,17 +397,6 @@ def get_face_crops(image_bgr, faces, target_width=500, target_height=130):
     return face_crops
 
 
-def get_face_crops2(image_bgr, faces):
-    # Convert color space from BGR to RGB since OpenCV uses BGR by default
-    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
-    face_crops = []
-    for face in faces:
-        # Extract the region of interest (the face) from the original image
-        x, y, w, h = face.left(), face.top(), face.width(), face.height()
-        face_crop = image_rgb[y:y+h, x:x+w]
-        face_pil = Image.fromarray(face_crop)
-        face_crops.append(face_pil)
-    return face_crops
 
 # Function to process uploaded images and display face crops
 def process_and_show_faces(image_input):
@@ -565,23 +433,24 @@ def process_selected_faces(image_input, selected_face_indices):
     for category, sticker_name in sticker_selections.items():
         if sticker_name:  # Check if a sticker was selected in this category
             # the sticker file path
-            sticker_path = os.path.join('stickers', category, sticker_name + '.png')
-
-            # Apply the selected sticker based on its category
-            if category == 'ears':
-                img_with_stickers = add_ears_sticker(img_with_stickers, sticker_path, faces)
-            elif category == 'glasses':
-                img_with_stickers = add_glasses_sticker(img_with_stickers, sticker_path, faces)
-            elif category == 'noses':
-                img_with_stickers = add_noses_sticker(img_with_stickers, sticker_path, faces)
-            elif category == 'headbands':
-                img_with_stickers = add_hats_sticker(img_with_stickers, sticker_path, faces)
-            elif category == 'hats':
-                img_with_stickers = add_hats_sticker(img_with_stickers, sticker_path, faces)
-            elif category == 'animal face':
-                img_with_stickers = add_animal_faces_sticker(img_with_stickers, sticker_path, faces)
-            else:
-                img_with_stickers = img_with_stickers
+            if sticker_name != 'None':
+                sticker_path = os.path.join('stickers', category, sticker_name + '.png')
+
+                # Apply the selected sticker based on its category
+                if category == 'ears':
+                    img_with_stickers = add_ears_sticker(img_with_stickers, sticker_path, faces)
+                elif category == 'glasses':
+                    img_with_stickers = add_glasses_sticker(img_with_stickers, sticker_path, faces)
+                elif category == 'noses':
+                    img_with_stickers = add_noses_sticker(img_with_stickers, sticker_path, faces)
+                elif category == 'headbands':
+                    img_with_stickers = add_hats_sticker(img_with_stickers, sticker_path, faces)
+                elif category == 'hats':
+                    img_with_stickers = add_hats_sticker(img_with_stickers, sticker_path, faces)
+                elif category == 'animal face':
+                    img_with_stickers = add_animal_faces_sticker(img_with_stickers, sticker_path, faces)
+                else:
+                    img_with_stickers = img_with_stickers
     # Convert back to PIL image
     img_with_stickers_pil = Image.fromarray(cv2.cvtColor(img_with_stickers, cv2.COLOR_BGR2RGB))
     
@@ -635,15 +504,44 @@ css = """
 }
 """
 
+def handle_image_upload(image):
+    global sticker_selections
+    sticker_selections = {category: "None" for category in STICKER_PATHS.keys()}
+    print("reset sticker_selections called") # Reset selections when a new image is loaded
+    # Print out the sticker selections state for each category
+    for category, selection in sticker_selections.items():
+        print(f"{category}: {selection}")
+    return image
+
+# Initialize the sticker selections dictionary
+def initialize_sticker_selections():
+    return {
+        'hats': None,
+        'animal face': None,
+        'ears': None,
+        'glasses': None,
+        'noses': None,
+        'headbands': None
+    }
+
+
+sticker_selections = initialize_sticker_selections()
+radio_components = {}
 # Create the Gradio interface
 with gr.Blocks(css=css) as demo:
     
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="pil", label="Original Image")
+            image_input.change(
+            handle_image_upload,
+            inputs=[image_input],
+            outputs=[image_input]
+            )
         with gr.Column():
             output_image = gr.Image(label="Image with Stickers")
         # Prepare the checkboxes and image placeholders
+
     detect_faces_btn = gr.Button("Detect Faces")
 
     with gr.Row():
@@ -661,7 +559,8 @@ with gr.Blocks(css=css) as demo:
     process_button = gr.Button("Apply Stickers To Selected Faces")
 
     process_button.click(
-        handle_face_selection, 
+        handle_face_selection,
+        # inputs=[image_input, face_checkboxes, sticker_selections], 
         inputs=[image_input] + face_checkboxes, 
         outputs=output_image
     )
@@ -682,20 +581,8 @@ with gr.Blocks(css=css) as demo:
                     choices = [sticker.split('/')[-1].replace('.png', '') for sticker in stickers]
                     radio = gr.Radio(label='', choices=choices, value="None", container=False, min_width=50, elem_classes="radio")
                     radio.change(lambda selection, cat=category: update_selections(cat, selection), inputs=[radio], outputs=[])
+                    radio_components[category] = radio  # Store the radio component
 
-    # function to clear all inputs and restart
-    def clear_all():
-        for i in range(MAX_EXPECTED_FACES):
-            face_images[i].value = None
-            face_checkboxes[i].value = False
-        output_image.value = None
-        for category in STICKER_TO_CATEGORY.values():
-            sticker_selections[category] = None
-            # radio.value = "None"
-            print(sticker_selections)
-                
-                
-    image_input.clear(fn=clear_all)
 
 demo.launch()