Refactor Gradio interface in app.py to enhance layout and usability, including updates to image input settings, background choice options, and button functionality for generating 3D shapes. Streamlined processing function connections for improved clarity.

app.py

@@ -237,97 +237,76 @@ _DESCRIPTION = '''
 '''
 
 # Create the Gradio interface
-with gr.Blocks(title="CRM: Single Image to 3D") 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():
             with gr.Row():
                 image_input = gr.Image(
                     label="Image input",
+                    image_mode="RGBA",
+                    sources="upload",
                     type="pil",
-                    height=300
-                )
-                processed_image = gr.Image(
-                    label="Processed Image",
-                    type="pil",
-                    height=300
                 )
-            
+                processed_image = gr.Image(label="Processed Image", interactive=False, type="pil", image_mode="RGB")
             with gr.Row():
                 with gr.Column():
-                    background_choice = gr.Radio(
-                        choices=["Alpha as mask", "Auto Remove background", "Custom Background"],
-                        value="Auto Remove background",
-                        label="Background choice"
-                    )
-                    back_ground_color = gr.ColorPicker(
-                        label="Background Color",
-                        value="#7F7F7F"
-                    )
+                    with gr.Row():
+                        background_choice = gr.Radio([
+                            "Alpha as mask",
+                            "Auto Remove background"
+                        ], value="Auto Remove background",
+                        label="Background choice")
+                    back_ground_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
+                        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", variant="primary")
-            
-            # Only add examples if the directory exists
+                    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")
             if os.path.exists("examples") and os.listdir("examples"):
                 gr.Examples(
-                    examples=[[os.path.join("examples", i)] for i in os.listdir("examples") if i.lower().endswith(('.png', '.jpg', '.jpeg'))],
-                    inputs=[image_input]
+                    examples=[os.path.join("examples", i) for i in os.listdir("examples") if i.lower().endswith(('.png', '.jpg', '.jpeg'))],
+                    inputs=[image_input],
+                    examples_per_page=20,
                 )
-        
         with gr.Column():
-            image_output = gr.Image(label="Output RGB image")
-            xyz_output = gr.Image(label="Output CCM image")
-            output_model = gr.Model3D(label="Output GLB")
+            image_output = gr.Image(interactive=False, label="Output RGB image")
+            xyz_output = gr.Image(interactive=False, label="Output CCM image")
+            output_model = gr.Model3D(
+                label="Output GLB",
+                interactive=False,
+            )
             gr.Markdown("Note: Ensure that the input image is correctly pre-processed into a grey background, otherwise the results will be unpredictable.")
 
-    # Connect the button to the processing function
-    text_button.click(
-        fn=process_and_generate,
-        inputs=[
-            image_input,
-            background_choice,
-            foreground_ratio,
-            back_ground_color,
-            seed,
-            guidance_scale,
-            step
-        ],
-        outputs=[
-            processed_image,
-            image_output,
-            xyz_output,
-            output_model
-        ],
-        api_name="predict"
+    inputs = [
+        processed_image,
+        seed,
+        guidance_scale,
+        step,
+    ]
+    outputs = [
+        image_output,
+        xyz_output,
+        output_model,
+    ]
+
+    text_button.click(fn=check_input_image, inputs=[image_input]).success(
+        fn=preprocess_image,
+        inputs=[image_input, background_choice, foreground_ratio, back_ground_color],
+        outputs=[processed_image],
+    ).success(
+        fn=gen_image,
+        inputs=inputs,
+        outputs=outputs,
     )
 
 if __name__ == "__main__":