|
- from tqdm import tqdm
- import cv2
- import numpy as np
- from numpy.linalg import norm
- from facex.model_zoo import get_model, compute_sim
- from facex.face_align import norm_crop, norm_crop2
-
- def decode_fourcc(v):
- v = int(v)
- return "".join([chr((v >> 8 * i) & 0xFF) for i in range(4)])
-
-
- class FSV:
- def __init__(self, ):
- self.face_detection = get_model("detection")
- self.face_detection.prepare(ctx_id=0, input_size=(640, 640), det_thresh=0.5)
- self.face_recognition = get_model("recognition")
- self.face_swapper = get_model("inswapper")
- self.fourcc_suffix = {
- "XVID": ".avi",
- "MP4V": ".MP4"
- }
-
- def get_face(self, image):
- _, kps = self.face_detection(image)
- face_image, M = norm_crop2(image, landmark=kps[0], image_size=512)
- # cv2.imwrite("target_face.jpg", face_image)
- return face_image
-
- def get_emb(self, face):
- embedding = self.face_recognition(face)
- norm_embedding = embedding / norm(embedding)
- return norm_embedding
-
- def swap(self, source, target):
- source_face, _ = self.get_face(source)
- target_face, M = self.get_face(target)
-
- norm_embedding = self.get_emb(source_face)
- swapped_face_image = self.face_swapper(target_face, norm_embedding)
- swapped_face_image = cv2.resize(swapped_face_image, (512, 512), interpolation=cv2.INTER_LANCZOS4)
- swapped_result = self.face_swapper.paste_back(swapped_face_image, source_face, target, M)
-
- return swapped_result
-
- def process_frame(self, frame, index):
- font = cv2.FONT_HERSHEY_SIMPLEX
-
- bboxes, kpses = self.face_detection(frame)
- for i, (bbox, kps) in enumerate(zip(bboxes, kpses)):
- x1, y1, x2, y2 = np.round(bbox[:4]).astype(int).tolist()
- face_image, M = norm_crop2(frame, landmark=kps, image_size=512)
- emb = self.get_emb(face_image)
- if compute_sim(emb, self.tgt_emb) < 0.5:
- continue
- swapped_face_image = self.face_swapper(face_image, self.src_emb)
- swapped_face_image = cv2.resize(swapped_face_image, (512, 512), interpolation=cv2.INTER_LANCZOS4)
- frame = self.face_swapper.paste_back(swapped_face_image, self.src_face, frame, M)
- # frame = cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 3, lineType=cv2.LINE_AA)
- # frame = cv2.putText(frame, "x", (x1, y1 - 5), font, 1.0, (0, 255, 0), 2)
- # cv2.imwrite(f"/root/swaped{index}.jpg", swapped_result)
-
- return frame
-
- def read_video(self, video_path):
- capture = cv2.VideoCapture(video_path)
- assert capture.isOpened(), f"open {video_path} failed!"
-
- fps = int(capture.get(cv2.CAP_PROP_FPS))
- resolution = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
- num_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
- fourcc = decode_fourcc(capture.get(cv2.CAP_PROP_FOURCC))
- print("fps:", fps)
- print("resolution:", resolution)
- print("num frames:", num_frames)
- print("fourcc: ", fourcc)
-
- cnt = 0
- frames = []
- pbar = tqdm(total=num_frames, ncols=100, desc="Read")
- while True:
- _, frame = capture.read()
- if frame is None:
- break
- cnt += 1
- pbar.update(1)
- frames.append(frame)
- pbar.close()
-
- capture.release()
-
- return frames, resolution, fps, fourcc
-
- def write_video(self, frames, video_path, resolution, fps, fourcc):
- writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*fourcc), fps, resolution)
- for frame in tqdm(frames, ncols=100, desc="Write"):
- writer.write(frame)
- writer.release()
-
- def __call__(self, src_image, tgt_image, video_path: str, save_path: str):
- src_image = cv2.imread(src_image)
- tgt_image = cv2.imread(tgt_image)
- self.src_face = self.get_face(src_image)
- self.src_emb = self.get_emb(self.src_face)
- self.tgt_emb = self.get_emb(self.get_face(tgt_image))
- frames, resolution, fps, fourcc = self.read_video(video_path)
- frames = [self.process_frame(frame, idx) for idx, frame in tqdm(enumerate(frames), total=len(frames), ncols=100, desc="Process")]
- fourcc = "XVID"
- suffix = self.fourcc_suffix[fourcc]
- self.write_video(frames, save_path+suffix, resolution, fps, fourcc)
-
- app = FSV()
- app("/root/8087885c969593fab7215dc46c14a59a.jpg", "/root/WangLiDan.jpg", "/root/WangLiDan.mp4", "/root/out")
|