Update infer.py

infer.py

@@ -16,6 +16,7 @@ from utils.seed_all import seed_all
 
 from contextlib import nullcontext
 import cv2
+from tqdm import tqdm  # 添加这一行以导入 tqdm
 
 check_min_version('0.28.0.dev0')
 
@@ -53,14 +54,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 = 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):
@@ -98,7 +99,7 @@ def lotus_video(input_video, task_name, seed, device):
     task_emb = torch.cat([torch.sin(task_emb), torch.cos(task_emb)], dim=-1).repeat(1, 1)
 
     output_g = []
-    for frame in frames:
+    for frame in tqdm(frames, desc="Processing frames"):  # 使用 tqdm 包裹 frames 列表
         if torch.backends.mps.is_available():
             autocast_ctx = nullcontext()
         else:
@@ -122,16 +123,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_npy = pred.mean(axis=-1)
                 # 修改为输出灰度图
-                output_color_g = Image.fromarray((output_npy_g * 255).astype(np.uint8), mode='L')
+                output_color = 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))
-                output_g.append(output_color_g)
- 
+
+            output_g.append(output_color_g)
 
     return output_g
 
@@ -307,14 +307,14 @@ def main():
 
             # Post-process the prediction
             save_file_name = os.path.basename(test_images[i])[:-4]
-            # 在 infer_pipe 函数中
-            if task_name == 'depth':
+            if args.task_name == 'depth':
                 output_npy = pred.mean(axis=-1)
                 # 修改为输出灰度图
                 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)