Update app.py

app.py

@@ -60,6 +60,7 @@ def process_image_check(path_input):
 def process_image(
     pipe,
     path_input,
+    mode='depth',
     processing_res=default_image_processing_res,
 ):
     name_base, name_ext = os.path.splitext(os.path.basename(path_input))
@@ -67,7 +68,6 @@ def process_image(
 
     path_output_dir = tempfile.mkdtemp()
     path_out_fp32 = os.path.join(path_output_dir, f"{name_base}_depth_fp32.npy")
-    path_out_16bit = os.path.join(path_output_dir, f"{name_base}_depth_16bit.png")
     path_out_vis = os.path.join(path_output_dir, f"{name_base}_depth_colored.png")
 
     input_image = Image.open(path_input)
@@ -81,11 +81,17 @@ def process_image(
 
     depth_pred = pipe_out.pred_np
     depth_colored = pipe_out.pred_colored
-    depth_16bit = (depth_pred * 65535.0).astype(np.uint16)
 
     np.save(path_out_fp32, depth_pred)
-    Image.fromarray(depth_16bit).save(path_out_16bit, mode="I;16")
+    
     depth_colored.save(path_out_vis)
+    
+    if mode == 'depth':
+        path_out_16bit = os.path.join(path_output_dir, f"{name_base}_depth_16bit.png")
+        depth_16bit = (depth_pred * 65535.0).astype(np.uint16)
+        Image.fromarray(depth_16bit).save(path_out_16bit, mode="I;16")
+    else:
+        path_out_16bit = None
 
     return (
         [path_out_16bit, path_out_vis],
@@ -93,7 +99,9 @@ def process_image(
     )
 
 def run_demo_server(pipe):
-    process_pipe_image = spaces.GPU(functools.partial(process_image, pipe))
+    process_pipe_depth = spaces.GPU(functools.partial(process_image, pipe, mode='depth'))
+    process_pipe_normal = spaces.GPU(functools.partial(process_image, pipe, mode='normal'))
+    process_pipe_dis = spaces.GPU(functools.partial(process_image, pipe, mode='dis'))
     gradio_theme = gr.themes.Default()
 
     with gr.Blocks(
@@ -164,7 +172,7 @@ def run_demo_server(pipe):
         )
 
         with gr.Tabs(elem_classes=["tabs"]):
-            with gr.Tab("Depth Estimation"):
+            with gr.Tab("Depth"):
                 with gr.Row():
                     with gr.Column():
                         image_input = gr.Image(
@@ -206,7 +214,7 @@ def run_demo_server(pipe):
                 filenames.extend(["line_%d.jpg" %(i+1) for i in range(6)])
                 filenames.extend(["real_%d.jpg" %(i+1) for i in range(24)])
                 Examples(
-                    fn=process_pipe_image,
+                    fn=process_pipe_depth,
                     examples=[
                         os.path.join("images", "depth", name)
                         for name in filenames
@@ -216,6 +224,109 @@ def run_demo_server(pipe):
                     cache_examples=True,
                     directory_name="examples_image",
                 )
+            
+            with gr.Tab("Normal"):
+                with gr.Row():
+                    with gr.Column():
+                        image_input = gr.Image(
+                            label="Input Image",
+                            type="filepath",
+                        )
+                        with gr.Row():
+                            image_submit_btn = gr.Button(
+                                value="Estimate Normal", variant="primary"
+                            )
+                            image_reset_btn = gr.Button(value="Reset")
+                        with gr.Accordion("Advanced options", open=False):
+                            image_processing_res = gr.Radio(
+                                [
+                                    ("Native", 0),
+                                    ("Recommended", 768),
+                                ],
+                                label="Processing resolution",
+                                value=default_image_processing_res,
+                            )
+                    with gr.Column():
+                        image_output_slider = ImageSlider(
+                            label="Predicted surface normal",
+                            type="filepath",
+                            show_download_button=True,
+                            show_share_button=True,
+                            interactive=False,
+                            elem_classes="slider",
+                            position=0.25,
+                        )
+                        image_output_files = gr.Files(
+                            label="Normal outputs",
+                            elem_id="download",
+                            interactive=False,
+                        )
+
+                filenames = []
+                filenames.extend(["%d.jpg" %(i+1) for i in range(10)])
+                Examples(
+                    fn=process_pipe_normal,
+                    examples=[
+                        os.path.join("images", "normal", name)
+                        for name in filenames
+                    ],
+                    inputs=[image_input],
+                    outputs=[image_output_slider, image_output_files],
+                    cache_examples=True,
+                    directory_name="examples_image",
+                )
+                
+            with gr.Tab("Dichotomous Segmentation"):
+                with gr.Row():
+                    with gr.Column():
+                        image_input = gr.Image(
+                            label="Input Image",
+                            type="filepath",
+                        )
+                        with gr.Row():
+                            image_submit_btn = gr.Button(
+                                value="Estimate Segmentation.", variant="primary"
+                            )
+                            image_reset_btn = gr.Button(value="Reset")
+                        with gr.Accordion("Advanced options", open=False):
+                            image_processing_res = gr.Radio(
+                                [
+                                    ("Native", 0),
+                                    ("Recommended", 768),
+                                ],
+                                label="Processing resolution",
+                                value=default_image_processing_res,
+                            )
+                    with gr.Column():
+                        image_output_slider = ImageSlider(
+                            label="Predicted dichotomous image segmentation",
+                            type="filepath",
+                            show_download_button=True,
+                            show_share_button=True,
+                            interactive=False,
+                            elem_classes="slider",
+                            position=0.25,
+                        )
+                        image_output_files = gr.Files(
+                            label="DIS outputs",
+                            elem_id="download",
+                            interactive=False,
+                        )
+
+                filenames = []
+                filenames.extend(["%d.jpg" %(i+1) for i in range(10)])
+                Examples(
+                    fn=process_pipe_dis,
+                    examples=[
+                        os.path.join("images", "dis", name)
+                        for name in filenames
+                    ],
+                    inputs=[image_input],
+                    outputs=[image_output_slider, image_output_files],
+                    cache_examples=True,
+                    directory_name="examples_image",
+                )
+
 
         ### Image tab
         image_submit_btn.click(