added mask creation

app.py

@@ -19,28 +19,36 @@ install_packages()
 import gradio as gr
 from control import main
 from diffusers.utils import load_image
+from mask import create_mask
 
-
-def process_images(image, mask, prompt):
+def process_image(image, prompt):
     try:
-        # Pass images directly to main function
+        # Create mask from input image
+        mask = create_mask(image)
+        
+        # First show the generated mask
+        yield mask
+        
+        # Then process image with mask
         result = main(image, mask, prompt)
-        return result
+        yield result
+        
     except Exception as e:
-        return str(e)
-
+        yield str(e)
 
 # Create Gradio interface
 demo = gr.Interface(
-    fn=process_images,
+    fn=process_image,
     inputs=[
         gr.Image(label="Input Image", type="pil"),
-        gr.Image(label="Mask Image", type="pil"),
         gr.Textbox(label="Prompt"),
     ],
-    outputs=gr.Image(label="Generated Image"),
+    outputs=[
+        gr.Image(label="Generated Mask"),
+        gr.Image(label="Generated Image")
+    ],
     title="Image Inpainting with FLUX ControlNet",
-    description="Upload an image and its mask, then provide a prompt to generate the inpainted result.",
+    description="Upload an image and provide a prompt. The system will first generate a mask and then create the inpainted result.",
 )
 
 if __name__ == "__main__":

mask.py

@@ -0,0 +1,45 @@
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+
+def create_mask(image):
+    """
+    Creates a binary mask from an input image by thresholding and smoothing.
+    
+    Args:
+        image: Input image (numpy array)
+        
+    Returns:
+        Binary mask image (numpy array)
+    """
+    # Store original image for visualization
+    image_org = image.copy()
+    
+    # Convert image to binary (0 or 255)
+    for i in range(len(image)):
+        for j in range(len(image[i])):
+            if image[i][j] != 255:
+                image[i][j] = 0
+                
+    # Add padding of 50 pixels on all sides
+    padding = 0
+    image = cv2.copyMakeBorder(image, padding, padding, padding, padding, 
+                              cv2.BORDER_CONSTANT, value=255)
+
+    # Apply Gaussian blur for smoothening
+    image = cv2.GaussianBlur(image, (5,5), 50)
+
+    # Threshold to create binary mask
+    _, mask = cv2.threshold(image, 254, 255, cv2.THRESH_BINARY)
+    
+    # Visualization (commented out for production use)
+    """
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
+    ax1.imshow(image_org, cmap="gray")
+    ax1.set_title("Original Image")
+    ax2.imshow(mask, cmap="gray")
+    ax2.set_title("Mask Image")
+    plt.show()
+    """
+    
+    return mask

requirements.txt

@@ -0,0 +1,4 @@
+diffusers==0.30.2
+torch>=2.6.0
+gradio>=5.17.1
+transformers>=4.49.0