Update inference.py

inference.py

@@ -247,33 +247,60 @@ class Predictor:
         # if is_gg_drive:
         #     temp_file = f'tmp_anime.{output_path.split(".")[-1]}'
 
-        def transform_and_save(self, frames, count):
-            transformed_frames = []
+        # 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)
+        #         transformed_frames.append(img)
+        #     return transformed_frames
+        def transform_and_write(frames, count, video_buffer):
             anime_images = self.transform(frames)
             for i in range(count):
                 img = np.clip(anime_images[i], 0, 255).astype(np.uint8)
-                transformed_frames.append(img)
-            return transformed_frames
+                success, encoded_image = cv2.imencode('.jpg', img)
+                if success:
+                    video_buffer.append(encoded_image.tobytes())
 
-        frame_count = len(video_frames)
-        transformed_video_frames = []
+        video_capture = cv2.VideoCapture(video_frames)
+        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))
 
-        batch_shape = (batch_size) + video_frames[0].shape
+        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
+
+        # frame_count = len(video_frames)
+        # transformed_video_frames = []
+         video_buffer = []
+
+        # batch_shape = (batch_size) + video_frames[0].shape
+        # frames = np.zeros(batch_shape, dtype=np.uint8)
+        # frame_idx = 0
+        _shape = (batch_size, frame_height, frame_width, 3)
         frames = np.zeros(batch_shape, dtype=np.uint8)
         frame_idx = 0
 
         try:
-            for frame in video_frames:
+            for _ in range(frame_count):
+                ret, frame = video_capture.read()
+                if not ret:
+                    break
                 frames[frame_idx] = frame
                 frame_idx += 1
                 if frame_idx == batch_size:
-                    transformed_frames = transform_and_save(frames, frame_idx)
-                    transformed_video_frames.extend(transformed_frames)
+                    transform_and_write(frames, frame_idx, video_buffer)
                     frame_idx = 0
         except Exception as e:
             print(e)
+        finally:
+            video_capture.release()
 
-        return transformed_video_frames
+        return video_buffer