facechain/face_adapter/face_preprocess.py at cfd5e26d837338d4276b6e69831074245d513381 · modelscope/facechain · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import cv2
import numpy as np
from skimage import transform as trans
from typing import List, Union, Optional

def read_image(img_path: str, mode: str = 'rgb', layout: str = 'HWC') -> np.ndarray:
    """
    Read and process the image based on the specified mode and layout.

    :param img_path: Path to the image file.
    :param mode: Image mode ('rgb' or 'gray').
    :param layout: Image layout ('HWC' or 'CHW').
    :return: Processed image.
    """
    if mode == 'gray':
        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    else:
        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
        if mode == 'rgb':
            img = img[..., ::-1]
        if layout == 'CHW':
            img = np.transpose(img, (2, 0, 1))
    return img

def preprocess(img: Union[str, np.ndarray], bbox: Optional[List[int]] = None, landmark: Optional[np.ndarray] = None,
               image_size: Optional[List[int]] = None, margin: int = 44, mode: str = 'rgb', layout: str = 'HWC') -> np.ndarray:
    """
    Preprocess the image by resizing and aligning based on the bounding box and landmark.

    :param img: Input image or image path.
    :param bbox: Bounding box coordinates.
    :param landmark: Landmark points.
    :param image_size: Desired image size.
    :param margin: Margin for bounding box.
    :param mode: Image mode ('rgb' or 'gray').
    :param layout: Image layout ('HWC' or 'CHW').
    :return: Preprocessed image.
    """
    if isinstance(img, str):
        img = read_image(img, mode=mode, layout=layout)

    image_size = image_size or [112, 112]

    if landmark is not None:
        src = np.array([
            [30.2946, 51.6963],
            [65.5318, 51.5014],
            [48.0252, 71.7366],
            [33.5493, 92.3655],
            [62.7299, 92.2041]
        ], dtype=np.float32)

        if image_size[1] == 112:
            src[:, 0] += 8.0

        dst = landmark.astype(np.float32)
        tform = trans.SimilarityTransform()
        tform.estimate(dst, src)
        M = tform.params[0:2, :]
    else:
        M = None

    if M is None:
        if bbox is None:
            det = np.zeros(4, dtype=np.int32)
            det[0] = int(img.shape[1] * 0.0625)
            det[1] = int(img.shape[0] * 0.0625)
            det[2] = img.shape[1] - det[0]
            det[3] = img.shape[0] - det[1]
        else:
            det = bbox

        bb = np.zeros(4, dtype=np.int32)
        bb[0] = np.maximum(det[0] - margin / 2, 0)
        bb[1] = np.maximum(det[1] - margin / 2, 0)
        bb[2] = np.minimum(det[2] + margin / 2, img.shape[1])
        bb[3] = np.minimum(det[3] + margin / 2, img.shape[0])
        ret = img[bb[1]:bb[3], bb[0]:bb[2], :]
        if len(image_size) > 0:
            ret = cv2.resize(ret, (image_size[1], image_size[0]))
        return ret
    else:
        warped = cv2.warpAffine(img, M, (image_size[1], image_size[0]), borderValue=0.0)
        return warped

def preprocess_3pt(img: Union[str, np.ndarray], bbox: Optional[List[int]] = None, landmark: Optional[np.ndarray] = None,
                   image_size: Optional[List[int]] = None, margin: int = 44, mode: str = 'rgb', layout: str = 'HWC') -> np.ndarray:
    """
    Preprocess the image by resizing and aligning based on the bounding box and landmark using 3 points.

    :param img: Input image or image path.
    :param bbox: Bounding box coordinates.
    :param landmark: Landmark points.
    :param image_size: Desired image size.
    :param margin: Margin for bounding box.
    :param mode: Image mode ('rgb' or 'gray').
    :param layout: Image layout ('HWC' or 'CHW').
    :return: Preprocessed image.
    """
    if isinstance(img, str):
        img = read_image(img, mode=mode, layout=layout)

    image_size = image_size or [112, 112]

    if landmark is not None:
        src = np.array([
            [30.2946, 51.6963],
            [65.5318, 51.5014],
            [48.0252, 71.7366]
        ], dtype=np.float32)

        if image_size[1] == 112:
            src[:, 0] += 8.0

        dst = landmark.astype(np.float32)
        tform = trans.SimilarityTransform()
        tform.estimate(dst, src)
        M = tform.params[0:2, :]
    else:
        M = None

    if M is None:
        if bbox is None:
            det = np.zeros(4, dtype=np.int32)
            det[0] = int(img.shape[1] * 0.0625)
            det[1] = int(img.shape[0] * 0.0625)
            det[2] = img.shape[1] - det[0]
            det[3] = img.shape[0] - det[1]
        else:
            det = bbox

        bb = np.zeros(4, dtype=np.int32)
        bb[0] = np.maximum(det[0] - margin / 2, 0)
        bb[1] = np.maximum(det[1] - margin / 2, 0)
        bb[2] = np.minimum(det[2] + margin / 2, img.shape[1])
        bb[3] = np.minimum(det[3] + margin / 2, img.shape[0])
        ret = img[bb[1]:bb[3], bb[0]:bb[2], :]
        if len(image_size) > 0:
            ret = cv2.resize(ret, (image_size[1], image_size[0]))
        return ret
    else:
        warped = cv2.warpAffine(img, M, (image_size[1], image_size[0]), borderValue=0.0)
        return warped