Update visualization/plot_3d_global.py

visualization/plot_3d_global.py

@@ -106,12 +106,13 @@ def plot_3d_motion(args, figsize=(10, 10), fps=120, radius=4):
         return frame
 
     out = []
-    for i in range(frame_number) : 
-        out.append(update(i))
-    out = np.stack(out, axis=0)
-    return torch.from_numpy(out)
-        
-
+    for i in range(frame_number):
+        frame = update(i)
+        if frame.ndim == 3 and frame.shape[2] == 3:  # Check that frame is H x W x 3
+            out.append(frame)
+        else:
+            print(f"Frame {i} has incorrect shape or channels: {frame.shape}")
+    return out
 
 def draw_to_batch(smpl_joints_batch, title_batch=None, outname=None):
     batch_size = len(smpl_joints_batch)