Update inference.py

inference.py

@@ -233,7 +233,70 @@ class Predictor:
         video_writer.close()
         #output_file.close()
 
-    def transform_video(self, video_frames, batch_size,start,end):
+
+        def transform_video(self, 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')
+
+        # 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]}'
+
+        def transform_and_write(frames, count, writer):
+            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 = cv2.VideoCapture(input_path)
+        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))
+
+        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)
+        frame_idx = 0
+
+        try:
+            for _ in tqdm(range(frame_count)):
+                ret, frame = video_capture.read()
+                if not ret:
+                    break
+                frames[frame_idx] = frame
+                frame_idx += 1
+                if frame_idx == batch_size:
+                    transform_and_write(frames, frame_idx, video_writer)
+                    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)
+
+        print(f'Animation video saved to {output_path}')    
+
+    def transform_video1(self, video, batch_size, start, end):
         #end = end or None
 
         # if not os.path.isfile(input_path):
@@ -262,11 +325,13 @@ class Predictor:
                 if success:
                     video_buffer.append(encoded_image.tobytes())
 
-        video_capture = cv2.VideoCapture(video_frames)
+        video_capture = cv2.VideoCapture(video)
         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))
+        
+        print(f'Transforming video {frame_count} frames, size: ({frame_width}, {frame_height})')
 
         if start or end:
             start_frame = int(start * fps)