YANGYYYY commited on
Commit
302545f
1 Parent(s): 0c394ef

Update inference.py

Browse files
Files changed (1) hide show
  1. inference.py +11 -11
inference.py CHANGED
@@ -234,19 +234,19 @@ class Predictor:
234
  #output_file.close()
235
 
236
 
237
- def transform_video(self, batch_size=4, start=0, end=0):
238
  end = end or None
239
 
240
- # if not os.path.isfile(input_path):
241
- # raise FileNotFoundError(f'{input_path} does not exist')
242
 
243
- # output_dir = "/".join(output_path.split("/")[:-1])
244
- # os.makedirs(output_dir, exist_ok=True)
245
- # is_gg_drive = '/drive/' in output_path
246
- # temp_file = ''
247
- #
248
- # if is_gg_drive:
249
- # temp_file = f'tmp_anime.{output_path.split(".")[-1]}'
250
 
251
  def transform_and_write(frames, count, writer):
252
  anime_images = self.transform(frames)
@@ -294,7 +294,7 @@ class Predictor:
294
  if temp_file:
295
  shutil.move(temp_file, output_path)
296
 
297
- print(f'Animation video saved to {output_path}')
298
 
299
  def transform_video1(self, video, batch_size, start, end):
300
  #end = end or None
 
234
  #output_file.close()
235
 
236
 
237
+ def transform_video(self, input_path, output_path, batch_size=4, start=0, end=0):
238
  end = end or None
239
 
240
+ if not os.path.isfile(input_path):
241
+ raise FileNotFoundError(f'{input_path} does not exist')
242
 
243
+ output_dir = "/".join(output_path.split("/")[:-1])
244
+ os.makedirs(output_dir, exist_ok=True)
245
+ is_gg_drive = '/drive/' in output_path
246
+ temp_file = ''
247
+
248
+ if is_gg_drive:
249
+ temp_file = f'tmp_anime.{output_path.split(".")[-1]}'
250
 
251
  def transform_and_write(frames, count, writer):
252
  anime_images = self.transform(frames)
 
294
  if temp_file:
295
  shutil.move(temp_file, output_path)
296
 
297
+ print(f'Animation video saved to {output_path}')
298
 
299
  def transform_video1(self, video, batch_size, start, end):
300
  #end = end or None