Refactor Gradio interface to improve layout and organization, enhancing user experience with clearer input/output sections and updated button functionality for generating 3D shapes.

app.py

@@ -206,90 +206,117 @@ _DESCRIPTION = '''
 * If you find the output unsatisfying, try using different seeds:)
 '''
 
-def generate(image, bg_choice, fg_ratio, bg_color, seed_val, guidance, steps):
-    if image is None:
-        raise gr.Error("No image uploaded!")
-    processed = preprocess_image(image, bg_choice, fg_ratio, bg_color)
-    return gen_image(processed, seed_val, guidance, steps)
-
-# Create a Blocks interface with minimal settings
-with gr.Blocks(
-    analytics_enabled=False,
-    title="CRM: Single Image to 3D Textured Mesh with Convolutional Reconstruction Model"
-) as demo:
+with gr.Blocks() as demo:
+    gr.Markdown("# CRM: Single Image to 3D Textured Mesh with Convolutional Reconstruction Model")
     gr.Markdown(_DESCRIPTION)
     
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(
-                label="Image input",
-                type="pil",
-                sources=["upload"]
-            )
-            bg_choice = gr.Radio(
-                choices=["Alpha as mask", "Auto Remove background"],
-                value="Auto Remove background",
-                label="Background choice"
-            )
-            fg_ratio = gr.Slider(
-                label="Foreground Ratio",
-                minimum=0.5,
-                maximum=1.0,
-                value=1.0,
-                step=0.05
+            with gr.Row():
+                image_input = gr.Image(
+                    label="Image input",
+                    image_mode="RGBA",
+                    sources=["upload"],
+                    type="pil"
+                )
+                processed_image = gr.Image(
+                    label="Processed Image",
+                    interactive=False,
+                    type="pil",
+                    image_mode="RGB"
+                )
+            
+            with gr.Row():
+                with gr.Column():
+                    with gr.Row():
+                        background_choice = gr.Radio(
+                            choices=["Alpha as mask", "Auto Remove background"],
+                            value="Auto Remove background",
+                            label="Background choice"
+                        )
+                    back_groud_color = gr.ColorPicker(
+                        label="Background Color",
+                        value="#7F7F7F",
+                        interactive=False
+                    )
+                    foreground_ratio = gr.Slider(
+                        label="Foreground Ratio",
+                        minimum=0.5,
+                        maximum=1.0,
+                        value=1.0,
+                        step=0.05
+                    )
+
+                with gr.Column():
+                    seed = gr.Number(
+                        value=1234,
+                        label="Seed",
+                        precision=0
+                    )
+                    guidance_scale = gr.Number(
+                        value=5.5,
+                        minimum=3,
+                        maximum=10,
+                        label="Guidance scale"
+                    )
+                    step = gr.Number(
+                        value=30,
+                        minimum=30,
+                        maximum=100,
+                        label="Sample steps",
+                        precision=0
+                    )
+            
+            text_button = gr.Button("Generate 3D shape")
+            
+            gr.Examples(
+                examples=[[os.path.join("examples", i)] for i in os.listdir("examples")],
+                inputs=[image_input],
+                examples_per_page=20,
+                cache_examples=False
             )
-            bg_color = gr.ColorPicker(
-                label="Background Color",
-                value="#7F7F7F",
+        
+        with gr.Column():
+            image_output = gr.Image(
+                label="Output RGB image",
                 interactive=False
             )
-            seed_val = gr.Number(
-                value=1234,
-                label="Seed",
-                precision=0
-            )
-            guidance = gr.Number(
-                value=5.5,
-                minimum=3,
-                maximum=10,
-                label="Guidance scale"
+            xyz_output = gr.Image(
+                label="Output CCM image",
+                interactive=False
             )
-            steps = gr.Number(
-                value=30,
-                minimum=30,
-                maximum=100,
-                label="Sample steps",
-                precision=0
+            output_model = gr.Model3D(
+                label="Output GLB",
+                interactive=False
             )
-            generate_btn = gr.Button("Generate")
-        
-        with gr.Column():
-            output_rgb = gr.Image(label="Output RGB image")
-            output_ccm = gr.Image(label="Output CCM image")
-            output_glb = gr.Model3D(label="Output GLB")
-    
-    # Connect the button click event
-    generate_btn.click(
-        fn=generate,
+            gr.Markdown("Note: Ensure that the input image is correctly pre-processed into a grey background, otherwise the results will be unpredictable.")
+
+    # Define the processing chain
+    def process_and_generate(image, bg_choice, fg_ratio, bg_color, seed_val, guidance, steps):
+        if image is None:
+            raise gr.Error("No image uploaded!")
+        processed = preprocess_image(image, bg_choice, fg_ratio, bg_color)
+        return gen_image(processed, seed_val, guidance, steps)
+
+    # Connect the button click event with API endpoint
+    text_button.click(
+        fn=process_and_generate,
         inputs=[
-            input_image,
-            bg_choice,
-            fg_ratio,
-            bg_color,
-            seed_val,
-            guidance,
-            steps
+            image_input,
+            background_choice,
+            foreground_ratio,
+            back_groud_color,
+            seed,
+            guidance_scale,
+            step
+        ],
+        outputs=[
+            image_output,
+            xyz_output,
+            output_model
         ],
-        outputs=[output_rgb, output_ccm, output_glb],
         api_name="predict"
     )
-    
-    # Add examples without caching
-    gr.Examples(
-        examples=[[os.path.join("examples", i)] for i in os.listdir("examples")],
-        inputs=input_image,
-        cache_examples=False
-    )
 
 if __name__ == "__main__":
-    demo.launch(show_api=True)
+    demo.queue().launch(show_api=True)