Update app.py

app.py

@@ -57,18 +57,17 @@ def process_image_check(path_input):
             "Missing image in the first pane: upload a file or use one from the gallery below."
         )
 
-def process_image(
+def process_depth(
     pipe,
     path_input,
-    mode='depth',
     processing_res=default_image_processing_res,
 ):
     name_base, name_ext = os.path.splitext(os.path.basename(path_input))
     print(f"Processing image {name_base}{name_ext}")
 
     path_output_dir = tempfile.mkdtemp()
-    path_out_fp32 = os.path.join(path_output_dir, f"{name_base}_%s_fp32.npy" %mode)
-    path_out_vis = os.path.join(path_output_dir, f"{name_base}_%s_colored.png" %mode)
+    path_out_fp32 = os.path.join(path_output_dir, f"{name_base}_depth_fp32.npy")
+    path_out_vis = os.path.join(path_output_dir, f"{name_base}_depth_colored.png")
 
     input_image = Image.open(path_input)
 
@@ -86,12 +85,79 @@ def process_image(
     np.save(path_out_fp32, depth_pred)
     depth_colored.save(path_out_vis)
     
-    if mode == 'depth':
-        path_out_16bit = os.path.join(path_output_dir, f"{name_base}_%s_16bit.png" %mode)
-        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
+    path_out_16bit = os.path.join(path_output_dir, f"{name_base}_%s_16bit.png" %mode)
+    depth_16bit = (depth_pred * 65535.0).astype(np.uint16)
+    Image.fromarray(depth_16bit).save(path_out_16bit, mode="I;16")
+
+    return (
+        [path_out_16bit, path_out_vis],
+        [path_out_16bit, path_out_fp32, path_out_vis],
+    )
+
+def process_normal(
+    pipe,
+    path_input,
+    processing_res=default_image_processing_res,
+):
+    name_base, name_ext = os.path.splitext(os.path.basename(path_input))
+    print(f"Processing image {name_base}{name_ext}")
+
+    path_output_dir = tempfile.mkdtemp()
+    path_out_fp32 = os.path.join(path_output_dir, f"{name_base}_normal_fp32.npy")
+    path_out_vis = os.path.join(path_output_dir, f"{name_base}_normal_colored.png")
+
+    input_image = Image.open(path_input)
+
+    pipe_out = pipe(
+        input_image,
+        processing_res=processing_res,
+        batch_size=1 if processing_res == 0 else 0,
+        show_progress_bar=False,
+        mode=mode,
+    )
+
+    depth_pred = pipe_out.pred_np
+    depth_colored = pipe_out.pred_colored
+
+    np.save(path_out_fp32, depth_pred)
+    depth_colored.save(path_out_vis)
+    
+    path_out_16bit = None
+
+    return (
+        [path_out_16bit, path_out_vis],
+        [path_out_16bit, path_out_fp32, path_out_vis],
+    )
+
+def process_dis(
+    pipe,
+    path_input,
+    processing_res=default_image_processing_res,
+):
+    name_base, name_ext = os.path.splitext(os.path.basename(path_input))
+    print(f"Processing image {name_base}{name_ext}")
+
+    path_output_dir = tempfile.mkdtemp()
+    path_out_fp32 = os.path.join(path_output_dir, f"{name_base}_dis_fp32.npy")
+    path_out_vis = os.path.join(path_output_dir, f"{name_base}_dis_colored.png")
+
+    input_image = Image.open(path_input)
+
+    pipe_out = pipe(
+        input_image,
+        processing_res=processing_res,
+        batch_size=1 if processing_res == 0 else 0,
+        show_progress_bar=False,
+        mode=mode,
+    )
+
+    depth_pred = pipe_out.pred_np
+    depth_colored = pipe_out.pred_colored
+
+    np.save(path_out_fp32, depth_pred)
+    depth_colored.save(path_out_vis)
+    
+    path_out_16bit = None
 
     return (
         [path_out_16bit, path_out_vis],
@@ -99,9 +165,9 @@ def process_image(
     )
 
 def run_demo_server(pipe_depth, pipe_normal, pipe_dis):
-    process_pipe_depth = spaces.GPU(functools.partial(process_image, pipe_depth, mode='depth'))
-    process_pipe_normal = spaces.GPU(functools.partial(process_image, pipe_normal, mode='normal'))
-    process_pipe_dis = spaces.GPU(functools.partial(process_image, pipe_dis, mode='seg'))
+    process_pipe_depth = spaces.GPU(functools.partial(process_depth, pipe_depth))
+    process_pipe_normal = spaces.GPU(functools.partial(process_normal, pipe_normal))
+    process_pipe_dis = spaces.GPU(functools.partial(process_dis, pipe_dis))
     gradio_theme = gr.themes.Default()
 
     with gr.Blocks(