jhj0517 commited on
Commit
b7be5e9
·
1 Parent(s): 04b83e7

Set frame rate and init dir paths

Browse files
modules/live_portrait/live_portrait_inferencer.py CHANGED
@@ -33,8 +33,17 @@ class LivePortraitInferencer:
33
  model_dir: str = MODELS_DIR,
34
  output_dir: str = OUTPUTS_DIR):
35
  self.model_dir = model_dir
36
- os.makedirs(os.path.join(self.model_dir, "animal"), exist_ok=True)
37
  self.output_dir = output_dir
 
 
 
 
 
 
 
 
 
 
38
  self.model_config = load_yaml(MODEL_CONFIG)["model_params"]
39
 
40
  self.appearance_feature_extractor = None
@@ -256,6 +265,8 @@ class LivePortraitInferencer:
256
  model_type=model_type
257
  )
258
 
 
 
259
  src_length = 1
260
 
261
  if src_image is not None:
@@ -268,7 +279,7 @@ class LivePortraitInferencer:
268
  else:
269
  self.psi_list = self.prepare_source(src_image, crop_factor)
270
 
271
- driving_images, vid_sound = extract_frames(driving_vid_path), extract_sound(driving_vid_path)
272
  driving_length = 0
273
  if driving_images is not None:
274
  if id(driving_images) != id(self.driving_images):
@@ -320,12 +331,12 @@ class LivePortraitInferencer:
320
  out = np.clip(psi.mask_ori * crop_with_fullsize + (1 - psi.mask_ori) * psi.src_rgb, 0, 255).astype(
321
  np.uint8)
322
 
323
- out_frame_path = get_auto_incremental_file_path(TEMP_VIDEO_OUT_FRAMES_DIR, "png")
324
  save_image(out, out_frame_path)
325
 
326
  progress(i/total_length, desc=f"Generating frames {i}/{total_length} ..")
327
 
328
- video_path = create_video_from_frames(TEMP_VIDEO_OUT_FRAMES_DIR)
329
 
330
  return video_path
331
 
 
33
  model_dir: str = MODELS_DIR,
34
  output_dir: str = OUTPUTS_DIR):
35
  self.model_dir = model_dir
 
36
  self.output_dir = output_dir
37
+ relative_dirs = [
38
+ os.path.join(self.model_dir, "animal"),
39
+ os.path.join(self.output_dir, "videos"),
40
+ os.path.join(self.output_dir, "temp"),
41
+ os.path.join(self.output_dir, "temp", "video_frames"),
42
+ os.path.join(self.output_dir, "temp", "video_frames", "out"),
43
+ ]
44
+ for dir_path in relative_dirs:
45
+ os.makedirs(dir_path, exist_ok=True)
46
+
47
  self.model_config = load_yaml(MODEL_CONFIG)["model_params"]
48
 
49
  self.appearance_feature_extractor = None
 
265
  model_type=model_type
266
  )
267
 
268
+ vid_info = get_video_info(vid_input=driving_vid_path)
269
+
270
  src_length = 1
271
 
272
  if src_image is not None:
 
279
  else:
280
  self.psi_list = self.prepare_source(src_image, crop_factor)
281
 
282
+ driving_images, vid_sound = extract_frames(driving_vid_path, os.path.join(self.output_dir, "temp", "video_frames")), extract_sound(driving_vid_path)
283
  driving_length = 0
284
  if driving_images is not None:
285
  if id(driving_images) != id(self.driving_images):
 
331
  out = np.clip(psi.mask_ori * crop_with_fullsize + (1 - psi.mask_ori) * psi.src_rgb, 0, 255).astype(
332
  np.uint8)
333
 
334
+ out_frame_path = get_auto_incremental_file_path(os.path.join(self.output_dir, "temp", "video_frames", "out"), "png")
335
  save_image(out, out_frame_path)
336
 
337
  progress(i/total_length, desc=f"Generating frames {i}/{total_length} ..")
338
 
339
+ video_path = create_video_from_frames(TEMP_VIDEO_OUT_FRAMES_DIR, frame_rate=vid_info.frame_rate, output_dir=os.path.join(self.output_dir, "videos"))
340
 
341
  return video_path
342