Update inference.py

inference.py

@@ -164,89 +164,21 @@ class Predictor:
             fname = fname.replace(f'.{ext}', '')
             cv2.imwrite(os.path.join(dest_dir, f'{fname}.jpg'), anime_img[..., ::-1])
 
-    def transform_video_yuan(self, input_path, output_path, batch_size=4, start=0, end=0):
-
-        '''
-        Transform a video to animation version
-        https://github.com/lengstrom/fast-style-transfer/blob/master/evaluate.py#L21
-        '''
-        # Force to None
-        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 = ''
-        #output_file = open(output_path, 'wb')
-
-        if is_gg_drive:
-            # Writing directly into google drive can be inefficient
-            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(0, count):
-                img = np.clip(anime_images[i], 0, 255)
-                writer.write_frame(img)
-
-        video_clip = VideoFileClip(input_path, audio=False)
-        if start or end:
-            video_clip = video_clip.subclip(start, end)
-
-        video_writer = ffmpeg_writer.FFMPEG_VideoWriter(
-            output_path,
-            video_clip.size, video_clip.fps, codec="libx264",
-            # preset="medium", bitrate="2000k",
-            audiofile=input_path, threads=None,
-            ffmpeg_params=None)
-
-        total_frames = round(video_clip.fps * video_clip.duration)
-        print(f'Transfroming video {input_path}, {total_frames} frames, size: {video_clip.size}')
-
-        batch_shape = (batch_size, video_clip.size[1], video_clip.size[0], 3)
-        frame_count = 0
-        frames = np.zeros(batch_shape, dtype=np.float32)
-        for frame in tqdm(video_clip.iter_frames()):
-            try:
-                frames[frame_count] = frame
-                frame_count += 1
-                if frame_count == batch_size:
-                    transform_and_write(frames, frame_count, video_writer)
-                    frame_count = 0
-            except Exception as e:
-                print(e)
-                break
-
-
-        # The last frames
-        if frame_count != 0:
-            transform_and_write(frames, frame_count, video_writer)
-
-        if temp_file:
-            # move to output path
-            shutil.move(temp_file, output_path)
-
-        print(f'Animation video saved to {output_path}')
-        video_writer.close()
-        #output_file.close()
 
 
     def transform_video(self, input_path, output_path, 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)
@@ -291,10 +223,10 @@ class Predictor:
             video_capture.release()
             video_writer.release()
 
-        if temp_file:
-            shutil.move(temp_file, output_path)
+        # 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 transform_video1(self, video, batch_size, start, end):
         #end = end or None