Update inference.py

inference.py

@@ -233,7 +233,7 @@ class Predictor:
         video_writer.close()
         #output_file.close()
 
-    def transform_video(self, input, batch_size=4, start=0, end=0):
+    def transform_video(self, video_frames, batch_size=4):
         end = end or None
 
         # if not os.path.isfile(input_path):
@@ -247,55 +247,37 @@ class Predictor:
         # if is_gg_drive:
         #     temp_file = f'tmp_anime.{output_path.split(".")[-1]}'
 
-        def transform_and_write(frames, count, writer):
+        def transform_and_save(self, frames, count):
+            transformed_frames = []
             anime_images = self.transform(frames)
             for i in range(count):
                 img = np.clip(anime_images[i], 0, 255).astype(np.uint8)
-                writer.write(img)
-
-        video_capture = input
-        frame_width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH))
-        frame_height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        fps = int(video_capture.get(cv2.CAP_PROP_FPS))
-        frame_count = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
-
-        if start or end:
-            start_frame = int(start * fps)
-            end_frame = int(end * fps) if end else frame_count
-            video_capture.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
-            frame_count = end_frame - start_frame
-
-        # video_writer = cv2.VideoWriter(
-        #     output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
-        video_buffer = []
+                transformed_frames.append(img)
+            return transformed_frames
 
-        #print(f'Transforming video {input_path}, {frame_count} frames, size: ({frame_width}, {frame_height})')
+        frame_count = len(video_frames)
+        transformed_video_frames = []
 
-        batch_shape = (batch_size, frame_height, frame_width, 3)
+        batch_shape = (batch_size,) + video_frames[0].shape
         frames = np.zeros(batch_shape, dtype=np.uint8)
         frame_idx = 0
 
         try:
-            for _ in tqdm(range(frame_count)):
-                ret, frame = video_capture.read()
-                if not ret:
-                    break
+            for frame in video_frames:
                 frames[frame_idx] = frame
                 frame_idx += 1
                 if frame_idx == batch_size:
-                    transform_and_write(frames, frame_idx, video_buffer)
+                    transformed_frames = self.transform_and_save(frames, frame_idx)
+                    transformed_video_frames.extend(transformed_frames)
                     frame_idx = 0
         except Exception as e:
             print(e)
-        finally:
-            video_capture.release()
-            #video_writer.release()
-            
-        return video_buffer
-        # if temp_file:
-        #     shutil.move(temp_file, output_path)
-
-        # print(f'Animation video saved to {output_path}')
+
+        return transformed_video_frames
+
+
+
+        
     def preprocess_images(self, images):
         '''
         Preprocess image for inference