Rollback changes

app.py

@@ -244,7 +244,7 @@ class ClothingInpainter:
 
         return images_output
 
-def process_segmentation(image, clothing_parts=None, prompt=None, dilation_iterations=2):
+def process_segmentation(image, dilation_iterations=2):
     try:
         if image is None:
             raise gr.Error("Please upload an image")
@@ -260,25 +260,13 @@ def process_segmentation(image, clothing_parts=None, prompt=None, dilation_itera
         if not all_masks:
             logger.error("No clothing detected in the image")
             raise gr.Error("No clothing detected in the image. Please try a different image.")
-            
         inpainter.last_mask = all_masks
-        
-        # If clothing parts are selected and prompt is provided, generate variations
-        if clothing_parts and prompt:
-            # Convert selected_parts to lowercase/dash format
-            selected_parts = [p.lower() for p in clothing_parts]
-            prompt_dict = {'pos': prompt}
-            
-            # Generate inpainted images
-            images = inpainter.inpaint(prompt_dict, image, selected_parts, dilation_iterations)
-            return images
-            
-        # Otherwise, just show the segmentation visualization
+        # Only show main clothing parts for selection
         main_parts = ['upper-clothes', 'dress', 'coat', 'pants', 'skirt']
         masks = {k: v for k, v in all_masks.items() if k in main_parts}
-        vis_image = inpainter.visualize_segmentation(image, masks, selected_parts=clothing_parts)
-        return [vis_image]
-        
+        vis_image = inpainter.visualize_segmentation(image, masks, selected_parts=None)
+        detected_parts = [k for k in masks.keys()]
+        return vis_image, gr.update(choices=detected_parts, value=[])
     except gr.Error as e:
         raise e
     except Exception as e:
@@ -362,110 +350,76 @@ def update_selected_parts(image, selected_parts, dilation_iterations):
 # Initialize the model
 init()
 
-def create_interface():
-    with gr.Blocks(title="ClothQuill - AI Clothing Inpainting", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# ClothQuill - AI Clothing Inpainting")
-        gr.Markdown("An AI-powered tool for seamlessly editing and modifying clothing in images using state-of-the-art deep learning models.")
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                # Input Section
-                with gr.Group():
-                    gr.Markdown("### Input Image")
-                    input_image = gr.Image(type="pil", label="Upload Image")
-                    with gr.Row():
-                        example_btn = gr.Button("Load Example Image", variant="secondary", size="sm")
-                        gr.Markdown("or upload your own image above")
-                    
-                    # Clothing Selection
-                    gr.Markdown("### Select Clothing Parts")
-                    clothing_parts = gr.CheckboxGroup(
-                        choices=[
-                            "upper-clothes", "dress", "coat", "pants", "skirt",
-                            "hat", "glove", "sunglasses", "scarf"
-                        ],
-                        label="Select parts to modify",
-                        value=["upper-clothes"]
-                    )
-                    
-                    # Advanced Options
-                    with gr.Accordion("Advanced Options", open=False):
-                        dilation_slider = gr.Slider(
-                            minimum=0,
-                            maximum=5,
-                            value=2,
-                            step=1,
-                            label="Mask Dilation (adjust for better blending)"
-                        )
-                
-                # Prompt Section
-                gr.Markdown("### Describe the New Clothing")
-                prompt = gr.Textbox(
-                    label="Enter your prompt",
-                    placeholder="e.g., A stylish black leather jacket with silver zippers",
-                    lines=2
-                )
-                
-                # Generate Button
-                generate_btn = gr.Button("Generate Variations", variant="primary")
-            
-            with gr.Column(scale=1):
-                # Output Section
-                gr.Markdown("### Results")
-                with gr.Group():
-                    output_gallery = gr.Gallery(
-                        label="Generated Variations",
-                        show_label=True,
-                        elem_id="gallery",
-                        columns=3,
-                        height="auto",
-                        object_fit="contain"
-                    )
-        
-        # Example Image Loading
-        def load_example():
-            example_path = "example.jpg"
-            if os.path.exists(example_path):
-                return Image.open(example_path)
-            return None
-        
-        example_btn.click(
-            fn=load_example,
-            outputs=input_image
-        )
-        
-        # Main Processing
-        generate_btn.click(
-            fn=process_segmentation,
-            inputs=[input_image, clothing_parts, prompt, dilation_slider],
-            outputs=output_gallery
-        )
-        
-        # Input Image Change Handler
-        input_image.change(
-            fn=process_segmentation,
-            inputs=[input_image, clothing_parts, prompt, dilation_slider],
-            outputs=output_gallery
-        )
-        
-        # Clothing Parts Change Handler
-        clothing_parts.change(
-            fn=process_segmentation,
-            inputs=[input_image, clothing_parts, prompt, dilation_slider],
-            outputs=output_gallery
-        )
+# Create Gradio interface
+with gr.Blocks(title="ClothQuill - AI Clothing Inpainting") as demo:
+    gr.Markdown("# ClothQuill - AI Clothing Inpainting")
+    gr.Markdown("Upload an image to see segmented clothing parts, then select parts to modify and describe your changes")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                type="pil", 
+                label="Upload Image",
+                scale=1,  # This ensures the image maintains its aspect ratio
+                height=None  # Allow dynamic height based on content
+            )
+            dilation_slider = gr.Slider(
+                minimum=0,
+                maximum=5,
+                value=2,
+                step=1,
+                label="Mask Dilation",
+                info="Adjust the mask dilation to control the area of modification"
+            )
+            selected_parts = gr.CheckboxGroup(
+                choices=[],
+                label="Select parts to modify",
+                value=[]
+            )
+            prompt = gr.Textbox(
+                label="Describe the clothing you want to generate",
+                placeholder="e.g., A stylish black leather jacket"
+            )
+            generate_btn = gr.Button("Generate")
         
-        # Dilation Slider Change Handler
-        dilation_slider.change(
-            fn=process_segmentation,
-            inputs=[input_image, clothing_parts, prompt, dilation_slider],
-            outputs=output_gallery
-        )
+        with gr.Column():
+            gallery = gr.Gallery(
+                label="Generated Results", 
+                show_label=False,
+                columns=2,
+                height=None,  # Allow dynamic height
+                object_fit="contain"  # Maintain aspect ratio
+            )
+    
+    # Add event handler for image upload
+    input_image.upload(
+        fn=process_segmentation,
+        inputs=[input_image, dilation_slider],
+        outputs=[input_image, selected_parts]
+    )
+    
+    # Add event handler for dilation changes
+    dilation_slider.change(
+        fn=update_dilation,
+        inputs=[input_image, selected_parts,dilation_slider],
+        outputs=input_image
+    )
     
-    return demo
+    # Add event handler for generation
+    generate_btn.click(
+        fn=process_image,
+        inputs=[prompt, input_image, selected_parts, dilation_slider],
+        outputs=gallery
+    )
+
+    # Add event handler for part selection changes
+    selected_parts.change(
+        fn=update_selected_parts,
+        inputs=[input_image, selected_parts, dilation_slider],
+        outputs=input_image
+    )
 
 if __name__ == "__main__":
-    demo = create_interface()
     demo.launch(share=True)