pred_color.py

###
'''
!git clone https://huggingface.co/spaces/radames/SPIGA-face-alignment-headpose-estimator
!cp -r SPIGA-face-alignment-headpose-estimator/SPIGA .
!pip install -r SPIGA/requirements.txt
!pip install datasets
!pip install retinaface-py>=0.0.2
!huggingface-cli login
'''

import sys
sys.path.insert(0, "SPIGA")

import numpy as np

from datasets import load_dataset
from spiga.inference.config import ModelConfig
from spiga.inference.framework import SPIGAFramework

processor = SPIGAFramework(ModelConfig("300wpublic"))

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)

    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)

import retinaface
import spiga.demo.analyze.track.retinasort.config as cfg

config = cfg.cfg_retinasort
device = "cpu"
face_detector = retinaface.RetinaFaceDetector(model=config['retina']['model_name'],
                                              device=device,
                                              extra_features=config['retina']['extra_features'],
                                              cfg_postreat=config['retina']['postreat'])

import cv2
Image = PIL.Image
import os
def single_pred_features(image):
    if type(image) == type("") and os.path.exists(image):
        image = Image.open(image).convert("RGB")
    elif hasattr(image, "shape"):
        image = Image.fromarray(image).convert("RGB")
    else:
        image = image.convert("RGB")
    image = image.resize((512, 512))
    cv2_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    face_detector.set_input_shape(image.size[1], image.size[0])
    features = face_detector.inference(image)
    if features:
        bboxes = features['bbox']
        bboxes_n = []
        for bbox in bboxes:
            x1, y1, x2, y2 = bbox[:4]
            bbox_wh = [x1, y1, x2-x1, y2-y1]
            bboxes_n.append(bbox_wh)
        face_features = processor.inference(cv2_image, bboxes_n)
        landmarks = face_features["landmarks"][0]
        face_features["spiga"] = landmarks
        face_features['spiga_seg'] = conditioning_from_landmarks(landmarks)
        return face_features

if __name__ == "__main__":
    from datasets import load_dataset, Dataset
    ds = load_dataset("svjack/facesyntheticsspigacaptioned_en_zh_1")
    dss = ds["train"]

    xiangbaobao = PIL.Image.open("babyxiang.png")
    out = single_pred_features(xiangbaobao.resize((512, 512)))
    out["spiga_seg"]

    out = single_pred_features(dss[0]["image"])
    out["spiga_seg"]