nicolasbuitragob commited on
Commit
0936f88
·
1 Parent(s): c903b19

rotate video

Browse files
Files changed (2) hide show
  1. tasks.py +5 -1
  2. vitpose.py +7 -5
tasks.py CHANGED
@@ -19,7 +19,11 @@ def process_video(video_path: str,vitpose: VitPose,user_id: str):
19
 
20
  vitpose.output_video_path = new_file_name
21
  annotated_frames = vitpose.run(video_path)
22
- annotated_video_path = vitpose.frames_to_video(annotated_frames,rotate=True)
 
 
 
 
23
 
24
  logger.info(f"Video processed {annotated_video_path}")
25
 
 
19
 
20
  vitpose.output_video_path = new_file_name
21
  annotated_frames = vitpose.run(video_path)
22
+
23
+
24
+
25
+
26
+ annotated_video_path = vitpose.frames_to_video(annotated_frames)
27
 
28
  logger.info(f"Video processed {annotated_video_path}")
29
 
vitpose.py CHANGED
@@ -81,16 +81,17 @@ class VitPose:
81
 
82
  return annotated_frame
83
 
84
- def frames_to_video(self, frames, rotate=False):
85
-
86
  fourcc = cv2.VideoWriter_fourcc(*'mp4v')
87
  height = self.video_metadata["height"]
88
  width = self.video_metadata["width"]
89
 
90
- # If rotation is requested, swap dimensions for the output video
 
 
 
91
  if rotate:
92
  print(f"Original dimensions: {width}x{height}, Rotated dimensions: {height}x{width}")
93
- # For the VideoWriter, we need to specify the dimensions of the output frames
94
  out = cv2.VideoWriter(self.output_video_path, fourcc, self.video_metadata["fps"], (height, width))
95
  else:
96
  print(f"Dimensions: {width}x{height}")
@@ -98,10 +99,11 @@ class VitPose:
98
 
99
  for frame in frames:
100
  if rotate:
101
- # Rotate 90 degrees clockwise
102
  rotated_frame = cv2.rotate(frame, cv2.ROTATE_90_COUNTERCLOCKWISE)
103
  out.write(rotated_frame)
104
  else:
 
105
  out.write(frame)
106
 
107
  out.release()
 
81
 
82
  return annotated_frame
83
 
84
+ def frames_to_video(self, frames):
 
85
  fourcc = cv2.VideoWriter_fourcc(*'mp4v')
86
  height = self.video_metadata["height"]
87
  width = self.video_metadata["width"]
88
 
89
+ # Always ensure vertical orientation
90
+ rotate = width > height # Rotate only if the video is in landscape mode
91
+
92
+ # For the VideoWriter, we need to specify the dimensions of the output frames
93
  if rotate:
94
  print(f"Original dimensions: {width}x{height}, Rotated dimensions: {height}x{width}")
 
95
  out = cv2.VideoWriter(self.output_video_path, fourcc, self.video_metadata["fps"], (height, width))
96
  else:
97
  print(f"Dimensions: {width}x{height}")
 
99
 
100
  for frame in frames:
101
  if rotate:
102
+ # Rotate landscape videos 90 degrees to make them vertical
103
  rotated_frame = cv2.rotate(frame, cv2.ROTATE_90_COUNTERCLOCKWISE)
104
  out.write(rotated_frame)
105
  else:
106
+ # Already vertical, no rotation needed
107
  out.write(frame)
108
 
109
  out.release()