Update infer.py

infer.py

@@ -53,13 +53,14 @@ def infer_pipe(pipe, image_input, task_name, seed, device):
             ).images[0]
 
         # Post-process the prediction
+        # 在 infer_pipe 函数中
         if task_name == 'depth':
             output_npy = pred.mean(axis=-1)
-            output_color = colorize_depth_map(output_npy)
+            # 修改为输出灰度图
+            output_color = Image.fromarray((output_npy * 255).astype(np.uint8), mode='L')
         else:
             output_npy = pred
             output_color = Image.fromarray((output_npy * 255).astype(np.uint8))
-
     return output_color
 
 def lotus_video(input_video, task_name, seed, device):
@@ -121,14 +122,15 @@ def lotus_video(input_video, task_name, seed, device):
                 task_emb=task_emb,
                 ).images[0]
             # Post-process the prediction
+            # 在 lotus_video 函数中
             if task_name == 'depth':
                 output_npy_g = pred_g.mean(axis=-1)
-                output_color_g = colorize_depth_map(output_npy_g)
+                # 修改为输出灰度图
+                output_color_g = Image.fromarray((output_npy_g * 255).astype(np.uint8), mode='L')
             else:
                 output_npy_g = pred_g
                 output_color_g = Image.fromarray((output_npy_g * 255).astype(np.uint8))
-
-            output_g.append(output_color_g)
+                output_g.append(output_color_g)
  
 
     return output_g
@@ -305,13 +307,14 @@ def main():
 
             # Post-process the prediction
             save_file_name = os.path.basename(test_images[i])[:-4]
-            if args.task_name == 'depth':
+            # 在 infer_pipe 函数中
+            if task_name == 'depth':
                 output_npy = pred.mean(axis=-1)
-                output_color = colorize_depth_map(output_npy)
+                # 修改为输出灰度图
+                output_color = Image.fromarray((output_npy * 255).astype(np.uint8), mode='L')
             else:
                 output_npy = pred
                 output_color = Image.fromarray((output_npy * 255).astype(np.uint8))
-
             output_color.save(os.path.join(output_dir_color, f'{save_file_name}.png'))
             np.save(os.path.join(output_dir_npy, f'{save_file_name}.npy'), output_npy)