Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import imageio | |
| import numpy as np | |
| from PIL import Image | |
| import torch | |
| torch.manual_seed(1024) | |
| from inference_utils import inference | |
| from face_utils import get_face_img, get_faces_video | |
| from batch_face import RetinaFace | |
| face_detector = RetinaFace(gpu_id=0) if torch.cuda.is_available() else RetinaFace(gpu_id=-1) | |
| def check_if_image_file(filename): | |
| return any(filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg", ".PNG", ".JPG", ".JPEG"]) | |
| def check_if_video_file(filename): | |
| return any(filename.endswith(extension) for extension in [".mp4", ".avi"]) | |
| def concat_image(image1, image2, image3): | |
| # resize to the same size of image3 | |
| image1 = image1.resize(image3.size) | |
| image2 = image2.resize(image3.size) | |
| concat_img = Image.new("RGB", (image3.width*3, image3.height)) | |
| concat_img.paste(image1, (0, 0)) | |
| concat_img.paste(image2, (image3.width, 0)) | |
| concat_img.paste(image3, (image3.width*2, 0)) | |
| return concat_img | |
| if __name__ == "__main__": | |
| import argparse | |
| from tqdm import tqdm | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--id_input", type=str, help="Path to the input, can be an image, a video", required=True) | |
| parser.add_argument("--makeup_reference", type=str, help="Path to the makeup image file", required=True) | |
| parser.add_argument("--fast_test", action="store_true", help="Use fast test mode, only process every 5 frames") | |
| parser.add_argument("--output_dir", type=str, default="./output") | |
| args = parser.parse_args() | |
| id_input = args.id_input | |
| makeup_reference = args.makeup_reference | |
| output_dir = args.output_dir | |
| os.makedirs(output_dir, exist_ok=True) | |
| # check if the input is a video or an image | |
| id_basename = os.path.basename(id_input).split(".")[0] | |
| if check_if_video_file(id_input): | |
| # read all frames from the video | |
| frames, coords = get_faces_video(face_detector, id_input) | |
| id_images = frames if not args.fast_test else frames[::5] | |
| coords = coords if not args.fast_test else coords[::5] | |
| elif check_if_image_file(id_input): | |
| frame, coord = get_face_img(face_detector, id_input) | |
| id_images = [frame] | |
| coords = [coord] | |
| else: | |
| raise ValueError("Unsupported file format for id_input") | |
| makeup_basename = os.path.basename(makeup_reference).split(".")[0] | |
| if check_if_image_file(makeup_reference): | |
| makeup_image_pil, _ = get_face_img(face_detector, makeup_reference) | |
| else: | |
| raise ValueError("Unsupported file format for makeup_reference") | |
| if len(id_images) == 0: | |
| raise ValueError("No input images loaded") | |
| elif len(id_images) == 1: | |
| result_img = inference(id_images[0], makeup_image_pil) | |
| # concat id, makeup and result images | |
| concat_img = concat_image(id_images[0], makeup_image_pil, result_img) | |
| concat_img.save(os.path.join(output_dir, id_basename + makeup_basename + '.png')) | |
| print(f"Output Image Saved to {os.path.join(output_dir, id_basename + makeup_basename + '.png')}") | |
| elif len(id_images) > 1: | |
| # get fps of the original video | |
| try: | |
| fps = imageio.get_reader(id_input).get_meta_data()["fps"] | |
| except: | |
| print("Failed to get the fps of the video, using default 25 fps") | |
| fps = 25 | |
| writer = imageio.get_writer(os.path.join(output_dir, id_basename + makeup_basename + '.mp4'), fps=fps if not args.fast_test else fps/5, quality=9, codec="libx264") | |
| for id_image_pil in tqdm(id_images): | |
| result_img = inference(id_image_pil, makeup_image_pil) | |
| concat_img = concat_image(id_image_pil, makeup_image_pil, result_img) | |
| writer.append_data(np.array(concat_img)) | |
| writer.close() | |
| print(f"Output Video Saved to {os.path.join(output_dir, id_basename + makeup_basename + '.mp4')}") | |