Update infer.py

infer.py

@@ -124,9 +124,9 @@ def lotus_video(input_video, task_name, seed, device):
                 ).images[0]
             # Post-process the prediction
             if task_name == 'depth':
-                output_npy = pred.mean(axis=-1)
+                output_npy = pred_g.mean(axis=-1)
                 # 修改为输出灰度图
-                output_color = Image.fromarray((output_npy * 255).astype(np.uint8), mode='L')
+                output_color_g = Image.fromarray((output_npy * 255).astype(np.uint8), mode='L')
             else:
                 output_npy_g = pred_g
                 output_color_g = Image.fromarray((output_npy_g * 255).astype(np.uint8))