upload

app.py

@@ -119,11 +119,11 @@ def process_normal(
         mode='normal',
     )
 
-    depth_pred = pipe_out.pred_np
-    depth_colored = pipe_out.pred_colored
+    normal_pred = pipe_out.pred_np
+    normal_colored = pipe_out.pred_colored
 
-    np.save(path_out_fp32, depth_pred)
-    depth_colored.save(path_out_vis)
+    np.save(path_out_fp32, normal_pred)
+    normal_colored.save(path_out_vis)
 
     return (
         [path_out_vis],
@@ -149,14 +149,14 @@ def process_dis(
         processing_res=processing_res,
         batch_size=1 if processing_res == 0 else 0,
         show_progress_bar=False,
-        mode='seg',
+        mode='dis',
     )
 
-    depth_pred = pipe_out.pred_np
-    depth_colored = pipe_out.pred_colored
+    dis_pred = pipe_out.pred_np
+    dis_colored = pipe_out.pred_colored
 
-    np.save(path_out_fp32, depth_pred)
-    depth_colored.save(path_out_vis)
+    np.save(path_out_fp32, dis_pred)
+    dis_colored.save(path_out_vis)
 
     return (
         [path_out_vis],
@@ -177,6 +177,39 @@ def process_matting(
 
     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='matting',
+    )
+
+    matting_pred = pipe_out.pred_np
+    matting_colored = pipe_out.pred_colored
+
+    np.save(path_out_fp32, matting_pred)
+    matting_colored.save(path_out_vis)
+
+    return (
+        [path_out_vis],
+        [path_out_fp32, path_out_vis],
+    )
+
+def process_seg(
+    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}_seg_fp32.npy")
+    path_out_vis = os.path.join(path_output_dir, f"{name_base}_seg_colored.png")
+
+    input_image = Image.open(path_input)
+
     pipe_out = pipe(
         input_image,
         processing_res=processing_res,
@@ -185,25 +218,64 @@ def process_matting(
         mode='seg',
     )
 
-    depth_pred = pipe_out.pred_np
-    depth_colored = pipe_out.pred_colored
+    seg_pred = pipe_out.pred_np
+    seg_colored = pipe_out.pred_colored
 
-    np.save(path_out_fp32, depth_pred)
-    depth_colored.save(path_out_vis)
+    np.save(path_out_fp32, seg_pred)
+    seg_colored.save(path_out_vis)
 
     return (
         [path_out_vis],
         [path_out_fp32, path_out_vis],
     )
 
+def process_disparity(
+    pipe,
+    path_input,
+    processing_res=default_image_processing_res,
+):
+    print('line 65', path_input)
+    
+    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}_disparity_fp32.npy")
+    path_out_vis = os.path.join(path_output_dir, f"{name_base}_disparity_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='disparity',
+    )
+
+    disparity_pred = pipe_out.pred_np
+    disparity_colored = pipe_out.pred_colored
+
+    np.save(path_out_fp32, disparity_pred)
+    disparity_colored.save(path_out_vis)
+    
+    path_out_16bit = os.path.join(path_output_dir, f"{name_base}_disparity_16bit.png")
+    disparity_16bit = (disparity_pred * 65535.0).astype(np.uint16)
+    Image.fromarray(disparity_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 run_demo_server(pipe_depth, pipe_normal, pipe_dis, pipe_matting, pipe_seg, pipe_disparity):
     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))
     process_pipe_matting = spaces.GPU(functools.partial(process_matting, pipe_matting))
-    process_pipe_seg = spaces.GPU(functools.partial(process_matting, pipe_seg))
-    process_pipe_disparity = spaces.GPU(functools.partial(process_matting, pipe_disparity))
+    process_pipe_seg = spaces.GPU(functools.partial(process_seg, pipe_seg))
+    process_pipe_disparity = spaces.GPU(functools.partial(process_disparity, pipe_disparity))
     gradio_theme = gr.themes.Default()
 
     with gr.Blocks(

genpercept/genpercept_pipeline.py

@@ -319,7 +319,7 @@ class GenPerceptPipeline(DiffusionPipeline):
             pred_colored_hwc = chw2hwc(pred_colored)
             pred_colored_img = Image.fromarray(pred_colored_hwc)
         else:
-            pred_colored_img = None
+            pred_colored_img = Image.fromarray((pipe_pred * 255.0).astype(np.uint8))
         
         if len(pipe_pred.shape) == 3 and pipe_pred.shape[0] == 3:
             pipe_pred = np.transpose(pipe_pred, (1, 2, 0))