Update inference.py

inference.py

@@ -233,19 +233,19 @@ class Predictor:
         video_writer.close()
         #output_file.close()
 
-    def transform_video(self, input_path, output_path, batch_size=4, start=0, end=0):
+    def transform_video(self, input, batch_size=4, start=0, end=0):
         end = end or None
 
-        if not os.path.isfile(input_path):
-            raise FileNotFoundError(f'{input_path} does not exist')
+        # if not os.path.isfile(input_path):
+        #     raise FileNotFoundError(f'{input_path} does not exist')
 
-        output_dir = "/".join(output_path.split("/")[:-1])
-        os.makedirs(output_dir, exist_ok=True)
-        is_gg_drive = '/drive/' in output_path
-        temp_file = ''
+        # output_dir = "/".join(output_path.split("/")[:-1])
+        # os.makedirs(output_dir, exist_ok=True)
+        # is_gg_drive = '/drive/' in output_path
+        # temp_file = ''
 
-        if is_gg_drive:
-            temp_file = f'tmp_anime.{output_path.split(".")[-1]}'
+        # if is_gg_drive:
+        #     temp_file = f'tmp_anime.{output_path.split(".")[-1]}'
 
         def transform_and_write(frames, count, writer):
             anime_images = self.transform(frames)
@@ -253,7 +253,7 @@ class Predictor:
                 img = np.clip(anime_images[i], 0, 255).astype(np.uint8)
                 writer.write(img)
 
-        video_capture = cv2.VideoCapture(input_path)
+        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))
@@ -265,10 +265,11 @@ class Predictor:
             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_writer = cv2.VideoWriter(
+        #     output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
+        video_buffer = []
 
-        print(f'Transforming video {input_path}, {frame_count} frames, size: ({frame_width}, {frame_height})')
+        #print(f'Transforming video {input_path}, {frame_count} frames, size: ({frame_width}, {frame_height})')
 
         batch_shape = (batch_size, frame_height, frame_width, 3)
         frames = np.zeros(batch_shape, dtype=np.uint8)
@@ -282,18 +283,19 @@ class Predictor:
                 frames[frame_idx] = frame
                 frame_idx += 1
                 if frame_idx == batch_size:
-                    transform_and_write(frames, frame_idx, video_writer)
+                    transform_and_write(frames, frame_idx, video_buffer)
                     frame_idx = 0
         except Exception as e:
             print(e)
         finally:
             video_capture.release()
-            video_writer.release()
-
-        if temp_file:
-            shutil.move(temp_file, output_path)
+            #video_writer.release()
+            
+        return video_buffer
+        # if temp_file:
+        #     shutil.move(temp_file, output_path)
 
-        print(f'Animation video saved to {output_path}')
+        # print(f'Animation video saved to {output_path}')
     def preprocess_images(self, images):
         '''
         Preprocess image for inference