padding_fix

demo_utils.py

@@ -39,7 +39,7 @@ def image_to_video(image_path, output_path, fps):
 
     img_array = np.array(img)
 
-    with imageio.get_writer(output_path, fps=fps, format='FFMPEG', codec='libx264', pixelformat='yuv420p') as writer:
+    with imageio.get_writer(output_path, fps=fps, format='FFMPEG', codec='libx264', pixelformat='yuv420p', macro_block_size=1) as writer:
         writer.append_data(img_array)
 
     return {
@@ -461,7 +461,10 @@ def process_video(video_path, setup, gazing_ratio=0.75, task_loss_requirement=0.
             else:
                 mask_tensor = mask_2d
 
-            mask_resized = F.interpolate(mask_tensor.unsqueeze(0).unsqueeze(0).float(), size=(scale_h, scale_w), mode='nearest')[0, 0]
+            # Map mask through padded space then crop to original image dimensions
+            H_pad, W_pad = nh * 224, nw * 224
+            mask_full = F.interpolate(mask_tensor.unsqueeze(0).unsqueeze(0).float(), size=(H_pad, W_pad), mode='nearest')[0, 0]
+            mask_resized = F.interpolate(mask_full[:H, :W].unsqueeze(0).unsqueeze(0), size=(scale_h, scale_w), mode='nearest')[0, 0]
 
             frame_tensor = video_viz[t]
             frame_scaled = F.interpolate(frame_tensor.unsqueeze(0), size=(scale_h, scale_w), mode='bicubic', align_corners=False).squeeze().clamp(0, 1)
@@ -503,7 +506,10 @@ def process_video(video_path, setup, gazing_ratio=0.75, task_loss_requirement=0.
             else:
                 mask_tensor_scale = mask_2d
 
-            mask_resized_scale = F.interpolate(mask_tensor_scale.unsqueeze(0).unsqueeze(0).float(), size=(scale_h, scale_w), mode='nearest')[0, 0]
+            # Map mask through padded space then crop to original image dimensions
+            H_pad, W_pad = nh * 224, nw * 224
+            mask_full_scale = F.interpolate(mask_tensor_scale.unsqueeze(0).unsqueeze(0).float(), size=(H_pad, W_pad), mode='nearest')[0, 0]
+            mask_resized_scale = F.interpolate(mask_full_scale[:H, :W].unsqueeze(0).unsqueeze(0), size=(scale_h, scale_w), mode='nearest')[0, 0]
 
             frame_tensor_scale = video_viz[t]
             frame_scaled_scale = F.interpolate(frame_tensor_scale.unsqueeze(0), size=(scale_h, scale_w), mode='bicubic', align_corners=False).squeeze().clamp(0, 1)
@@ -576,6 +582,6 @@ def process_video(video_path, setup, gazing_ratio=0.75, task_loss_requirement=0.
 
 
 def save_video(frames, output_path, fps):
-    with imageio.get_writer(output_path, fps=fps, format='FFMPEG', codec='libx264', pixelformat='yuv420p') as writer:
+    with imageio.get_writer(output_path, fps=fps, format='FFMPEG', codec='libx264', pixelformat='yuv420p', macro_block_size=1) as writer:
         for frame in frames:
             writer.append_data(frame)