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(
|
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 |
|