import os import cv2 import tqdm import numpy as np from PIL import Image from facelib import FaceDetector from spiga.inference.config import ModelConfig from spiga.inference.framework import SPIGAFramework processor = SPIGAFramework(ModelConfig("300wpublic")) def center_crop(image, size): width, height = image.size left = (width - size) // 2 top = (height - size) // 2 right = left + size bottom = top + size cropped_image = image.crop((left, top, right, bottom)) return cropped_image def resize(image, size): width, height = image.size if width > height: # 按宽度进行比例调整 new_width = size new_height = int(height * (size / width)) else: # 按高度进行比例调整 new_height = size new_width = int(width * (size / height)) resized_image = image.resize((new_width, new_height)) return resized_image def preprocess(example, name, path): image = resize(example, 512) # 调用中心剪裁函数 cropped_image = center_crop(image, 512) # 保存剪裁后的图像 cropped_image.save(path+name) return cropped_image # We obtain the bbox from the existing landmarks in the dataset. # We could use `dlib`, but this should be faster. # Note that the `landmarks` are stored as strings. def get_landmarks(image, detector): image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR) faces, boxes, scores, landmarks = detector.detect_align(image) # 一定要用align啊 boxes = boxes.cpu().numpy() box_ls = [] for box in boxes: x, y, x1, y1 = box box = x, y, x1 - x, y1 - y box_ls.append(box) if len(box_ls) == 0: return [] else: features = processor.inference(image, box_ls) landmarks = np.array(features['landmarks']) return landmarks def parse_landmarks(landmarks): ldm = [] for landmark in landmarks: ldm.append([(float(x), float(y)) for x, y in landmark]) return ldm def bbox_from_landmarks(landmarks_): landmarks = parse_landmarks(landmarks_) bbox = [] for ldm in landmarks: landmarks_x, landmarks_y = zip(*ldm) x_min, x_max = min(landmarks_x), max(landmarks_x) y_min, y_max = min(landmarks_y), max(landmarks_y) width = x_max - x_min height = y_max - y_min # Give it a little room; I think it works anyway x_min -= 5 y_min -= 5 width += 10 height += 10 bbox.append((x_min, y_min, width, height)) return bbox def spiga_process(example, detector): ldms = get_landmarks(example, detector) if len(ldms) == 0: return False else: image = example image = np.array(image) # BGR image = image[:, :, ::-1] bbox = bbox_from_landmarks(ldms) features = processor.inference(image, [*bbox]) landmarks = features["landmarks"] spigas = landmarks return spigas # For some reason this map doesn't work with num_proc > 1 :( # TODO: run inference on GPU # ## "Segmentation" # We use bezier paths to draw contours and areas. import matplotlib.pyplot as plt import matplotlib.patches as patches from matplotlib.path import Path import PIL def get_patch(landmarks, color='lime', closed=False): contour = landmarks ops = [Path.MOVETO] + [Path.LINETO] * (len(contour) - 1) facecolor = (0, 0, 0, 0) # Transparent fill color, if open if closed: contour.append(contour[0]) ops.append(Path.CLOSEPOLY) facecolor = color path = Path(contour, ops) return patches.PathPatch(path, facecolor=facecolor, edgecolor=color, lw=4) # Draw to a buffer. def conditioning_from_landmarks(landmarks_, size=512): # Precisely control output image size dpi = 72 fig, ax = plt.subplots(1, figsize=[size / dpi, size / dpi], tight_layout={'pad': 0}) fig.set_dpi(dpi) black = np.zeros((size, size, 3)) ax.imshow(black) for landmarks in landmarks_: face_patch = get_patch(landmarks[0:17]) l_eyebrow = get_patch(landmarks[17:22], color='yellow') r_eyebrow = get_patch(landmarks[22:27], color='yellow') nose_v = get_patch(landmarks[27:31], color='orange') nose_h = get_patch(landmarks[31:36], color='orange') l_eye = get_patch(landmarks[36:42], color='magenta', closed=True) r_eye = get_patch(landmarks[42:48], color='magenta', closed=True) outer_lips = get_patch(landmarks[48:60], color='cyan', closed=True) inner_lips = get_patch(landmarks[60:68], color='blue', closed=True) ax.add_patch(face_patch) ax.add_patch(l_eyebrow) ax.add_patch(r_eyebrow) ax.add_patch(nose_v) ax.add_patch(nose_h) ax.add_patch(l_eye) ax.add_patch(r_eye) ax.add_patch(outer_lips) ax.add_patch(inner_lips) plt.axis('off') fig.canvas.draw() buffer, (width, height) = fig.canvas.print_to_buffer() assert width == height assert width == size buffer = np.frombuffer(buffer, np.uint8).reshape((height, width, 4)) buffer = buffer[:, :, 0:3] plt.close(fig) return PIL.Image.fromarray(buffer) def spiga_segmentation(spiga, size): landmarks = spiga spiga_seg = conditioning_from_landmarks(landmarks, size=size) return spiga_seg if __name__ == '__main__': # ## Obtain SPIGA features processor = SPIGAFramework(ModelConfig("300wpublic")) detector = FaceDetector(weight_path="/share2/zhangyuxuan/project/train_ip_cn/datasets/make_kps/pretrained_models/mobilenet0.25_Final.pth") id_folder = "/share2/zhangyuxuan/project/train_ip_cn/test_img_2/id/" pose_folder = "/share2/zhangyuxuan/project/train_ip_cn/test_img_2/pose/" if not os.path.exists(pose_folder): os.makedirs(pose_folder) pbar = tqdm.tqdm(os.listdir(id_folder)) for name in pbar: face = Image.open(id_folder+name).convert("RGB").resize((512, 512)) face.save(id_folder+name) spigas = spiga_process(face, detector) if spigas == False: height = 512 width = 512 channels = 3 black_image = np.zeros((height, width, channels), dtype=np.uint8) black_image_cv2 = cv2.cvtColor(black_image, cv2.COLOR_RGB2BGR) continue else: spigas_faces = spiga_segmentation(spigas) spigas_faces.save(pose_folder + name)