garmentiq.landmark.detection.utils

  1import math
  2import numpy as np
  3import cv2
  4from PIL import Image
  5import torchvision.transforms as transforms
  6from typing import Union
  7
  8
  9def get_max_preds(batch_heatmaps):
 10    """
 11    get predictions from score maps
 12    heatmaps: numpy.ndarray([batch_size, num_joints, height, width])
 13
 14    Args:
 15        batch_heatmaps (numpy.ndarray): Heatmaps generated by the model.
 16                                       Shape: [batch_size, num_joints, height, width].
 17
 18    Returns:
 19        tuple:
 20            - preds (numpy.ndarray): Predicted coordinates.
 21            - maxvals (numpy.ndarray): Maximum values (confidence scores) for each prediction.
 22    """
 23    assert isinstance(
 24        batch_heatmaps, np.ndarray
 25    ), "batch_heatmaps should be numpy.ndarray"
 26    assert batch_heatmaps.ndim == 4, "batch_images should be 4-ndim"
 27
 28    batch_size = batch_heatmaps.shape[0]
 29    num_joints = batch_heatmaps.shape[1]
 30    width = batch_heatmaps.shape[3]
 31    heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1))
 32    idx = np.argmax(heatmaps_reshaped, 2)
 33    maxvals = np.amax(heatmaps_reshaped, 2)
 34
 35    maxvals = maxvals.reshape((batch_size, num_joints, 1))
 36    idx = idx.reshape((batch_size, num_joints, 1))
 37
 38    preds = np.tile(idx, (1, 1, 2)).astype(np.float32)
 39
 40    preds[:, :, 0] = (preds[:, :, 0]) % width
 41    preds[:, :, 1] = np.floor((preds[:, :, 1]) / width)
 42
 43    pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2))
 44    pred_mask = pred_mask.astype(np.float32)
 45
 46    preds *= pred_mask
 47    return preds, maxvals
 48
 49
 50def get_final_preds(output, height=96, width=72):
 51    """
 52    Transforms raw heatmap outputs into final landmark coordinates.
 53
 54    Applies post-processing (e.g., quarter offset for sub-pixel accuracy)
 55    to refine the landmark predictions from heatmaps.
 56
 57    Args:
 58        output (numpy.ndarray): Raw heatmap output from the model.
 59        height (int): Height of the heatmap. Defaults to 96.
 60        width (int): Width of the heatmap. Defaults to 72.
 61
 62    Returns:
 63        tuple:
 64            - coords (numpy.ndarray): Final predicted coordinates.
 65            - maxvals (numpy.ndarray): Confidence scores for the predictions.
 66    """
 67    heatmap_height = height
 68    heatmap_width = width
 69
 70    batch_heatmaps = output
 71    coords, maxvals = get_max_preds(batch_heatmaps)
 72    # post-processing
 73    for n in range(coords.shape[0]):
 74        for p in range(coords.shape[1]):
 75            hm = batch_heatmaps[n][p]
 76            px = int(math.floor(coords[n][p][0] + 0.5))
 77            py = int(math.floor(coords[n][p][1] + 0.5))
 78            if 1 < px < heatmap_width - 1 and 1 < py < heatmap_height - 1:
 79                diff = np.array(
 80                    [hm[py][px + 1] - hm[py][px - 1], hm[py + 1][px] - hm[py - 1][px]]
 81                )
 82                coords[n][p] += np.sign(diff) * 0.25
 83
 84    return coords, maxvals
 85
 86
 87def flip_back(output_flipped, matched_parts, heatmap_wid):
 88    """
 89    Flips the output (coordinates or heatmaps) horizontally for test-time augmentation.
 90
 91    Args:
 92        output_flipped (numpy.ndarray): The output (heatmaps or coordinates) that has been flipped.
 93        matched_parts (list): A list of tuples indicating which joint pairs are left-right symmetric.
 94        heatmap_wid (int): The width of the heatmap (used for coordinate flipping).
 95
 96    Returns:
 97        numpy.ndarray: The flipped output with joints correctly reordered.
 98    """
 99    if output_flipped.ndim == 4:
100        output_flipped = output_flipped[:, :, :, ::-1]
101        for pair in matched_parts:
102            tmp = output_flipped[:, pair[0], :, :].copy()
103            output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
104            output_flipped[:, pair[1], :, :] = tmp
105    elif output_flipped.ndim == 3:
106        output_flipped[:, :, 0] = heatmap_wid - output_flipped[:, :, 0]
107        for pair in matched_parts:
108            tmp = output_flipped[:, pair[0], :].copy()
109            output_flipped[:, pair[0], :] = output_flipped[:, pair[1], :]
110            output_flipped[:, pair[1], :] = tmp
111    else:
112        raise NotImplementedError(
113            "output_flipped should be [batch_size, num_joints, height, width], "
114            "or [batch_size, num_joints, coord_dim"
115        )
116
117    return output_flipped
118
119
120def fliplr_joints(joints, joints_vis, width, matched_parts):
121    """
122    Flips joint coordinates horizontally and reorders them based on matched parts.
123
124    Args:
125        joints (numpy.ndarray): Array of joint coordinates.
126        joints_vis (numpy.ndarray): Array indicating visibility of joints.
127        width (int): Width of the image or feature map.
128        matched_parts (list): A list of tuples indicating which joint pairs are left-right symmetric.
129
130    Returns:
131        tuple:
132            - joints (numpy.ndarray): Flipped and reordered joint coordinates.
133            - joints_vis (numpy.ndarray): Corresponding joint visibility.
134    """
135    # Flip horizontal
136    joints[:, 0] = width - joints[:, 0] - 1
137
138    # Change left-right parts
139    for pair in matched_parts:
140        joints[pair[0], :], joints[pair[1], :] = (
141            joints[pair[1], :],
142            joints[pair[0], :].copy(),
143        )
144        joints_vis[pair[0], :], joints_vis[pair[1], :] = (
145            joints_vis[pair[1], :],
146            joints_vis[pair[0], :].copy(),
147        )
148
149    return joints * joints_vis, joints_vis
150
151
152def transform_preds(coords, center, scale, output_size: list[int, int] = [72, 96]):
153    """
154    Transforms predicted coordinates from heatmap space back to original image space.
155
156    Args:
157        coords (numpy.ndarray): Predicted coordinates in heatmap space.
158        center (numpy.ndarray): Center of the original image (or cropped region).
159        scale (numpy.ndarray): Scale factor applied during preprocessing.
160        output_size (list[int, int], optional): The size of the output image after transformation. Defaults to [72, 96].
161
162    Returns:
163        numpy.ndarray: Transformed coordinates in original image space.
164    """
165    target_coords = np.zeros(coords.shape)
166    trans = get_affine_transform(center, scale, 0, output_size, inv=1)
167    for p in range(coords.shape[0]):
168        target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
169    return target_coords
170
171
172def get_affine_transform(
173    center, scale, rot, output_size, shift=np.array([0, 0], dtype=np.float32), inv=0
174):
175    """
176    Calculates the 2x3 affine transformation matrix for image cropping and resizing.
177
178    Args:
179        center (numpy.ndarray): Center of the original image or region of interest.
180        scale (numpy.ndarray): Scale factor for the transformation.
181        rot (float): Rotation angle in degrees.
182        output_size (list): Target output size [width, height].
183        shift (numpy.ndarray, optional): Shift applied to the center. Defaults to [0, 0].
184        inv (int, optional): If 1, returns the inverse transformation matrix. Defaults to 0.
185
186    Returns:
187        numpy.ndarray: The 2x3 affine transformation matrix.
188    """
189    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
190        print(scale)
191        scale = np.array([scale, scale])
192
193    scale_tmp = scale * 200.0
194    src_w = scale_tmp[0]
195    dst_w = output_size[0]
196    dst_h = output_size[1]
197
198    rot_rad = np.pi * rot / 180
199    src_dir = get_dir([0, src_w * -0.5], rot_rad)
200    dst_dir = np.array([0, dst_w * -0.5], np.float32)
201
202    src = np.zeros((3, 2), dtype=np.float32)
203    dst = np.zeros((3, 2), dtype=np.float32)
204    src[0, :] = center + scale_tmp * shift
205    src[1, :] = center + src_dir + scale_tmp * shift
206    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
207    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
208
209    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
210    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
211
212    if inv:
213        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
214    else:
215        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
216
217    return trans
218
219
220def affine_transform(pt, t):
221    """
222    Applies an affine transformation matrix to a 2D point.
223
224    Args:
225        pt (tuple or list): The 2D point (x, y) to transform.
226        t (numpy.ndarray): The 2x3 affine transformation matrix.
227
228    Returns:
229        numpy.ndarray: The transformed 2D point.
230    """
231    new_pt = np.array([pt[0], pt[1], 1.0]).T
232    new_pt = np.dot(t, new_pt)
233    return new_pt[:2]
234
235
236def get_3rd_point(a, b):
237    """
238    Calculates a third point to form a right-angled triangle with two given points.
239    Used for creating a 3-point basis for affine transformations.
240
241    Args:
242        a (numpy.ndarray): First point.
243        b (numpy.ndarray): Second point.
244
245    Returns:
246        numpy.ndarray: The calculated third point.
247    """
248    direct = a - b
249    return b + np.array([-direct[1], direct[0]], dtype=np.float32)
250
251
252def get_dir(src_point, rot_rad):
253    """
254    Calculates the direction vector after rotation.
255
256    Args:
257        src_point (list): Source point [x, y].
258        rot_rad (float): Rotation angle in radians.
259
260    Returns:
261        list: The rotated direction vector.
262    """
263    sn, cs = np.sin(rot_rad), np.cos(rot_rad)
264
265    src_result = [0, 0]
266    src_result[0] = src_point[0] * cs - src_point[1] * sn
267    src_result[1] = src_point[0] * sn + src_point[1] * cs
268
269    return src_result
270
271
272def crop(img, center, scale, output_size, rot=0):
273    """
274    Crops and resizes an image using an affine transformation.
275
276    Args:
277        img (numpy.ndarray): The input image.
278        center (numpy.ndarray): The center of the crop region.
279        scale (numpy.ndarray): The scale factor for the crop.
280        output_size (tuple): The target output size (width, height).
281        rot (int, optional): Rotation angle in degrees. Defaults to 0.
282
283    Returns:
284        numpy.ndarray: The cropped and transformed image.
285    """
286    trans = get_affine_transform(center, scale, rot, output_size)
287
288    dst_img = cv2.warpAffine(
289        img, trans, (int(output_size[0]), int(output_size[1])), flags=cv2.INTER_LINEAR
290    )
291
292    return dst_img
293
294
295def input_image_transform(
296    img_input: Union[str, np.ndarray],
297    scale_std: float = 200.0,
298    resize_dim: list[int] = [288, 384],
299    normalize_mean: list[float] = [0.485, 0.456, 0.406],
300    normalize_std: list[float] = [0.229, 0.224, 0.225],
301):
302    """
303    Preprocesses an input image for landmark detection.
304
305    This function takes an image (either path or NumPy array), applies an affine
306    transformation (scaling, centering), resizes it, converts it to a PyTorch tensor,
307    and normalizes it.
308
309    Args:
310        img_input (Union[str, np.ndarray]): Path to the image file or a NumPy array of the image.
311        scale_std (float, optional): Standard scale for image transformation. Defaults to 200.0.
312        resize_dim (list[int], optional): Target dimensions [width, height] for the transformed image.
313                                         Defaults to [288, 384].
314        normalize_mean (list[float], optional): Mean values for image normalization (RGB channels).
315                                                Defaults to [0.485, 0.456, 0.406].
316        normalize_std (list[float], optional): Standard deviation values for image normalization (RGB channels).
317                                               Defaults to [0.229, 0.224, 0.225].
318
319    Raises:
320        ValueError: If `img_input` is neither a file path nor a NumPy array.
321
322    Returns:
323        tuple:
324            - input_tensor (torch.Tensor): The preprocessed image as a PyTorch tensor, ready for model input.
325            - image_np (numpy.ndarray): The original image as a NumPy array (RGB).
326            - center (numpy.ndarray): The center of the original image used for transformation.
327            - scale (numpy.ndarray): The scale factor used for transformation.
328    """
329    if isinstance(img_input, str):
330        img = Image.open(img_input).convert("RGB")
331    elif isinstance(img_input, np.ndarray):
332        img = Image.fromarray(img_input.astype(np.uint8))
333    else:
334        raise ValueError("img_input must be a file path or a NumPy array.")
335
336    image_np = np.array(img)
337
338    h, w = image_np.shape[:2]
339    center = np.array([w / 2, h / 2], dtype=np.float32)
340    scale = np.array([w / scale_std, h / scale_std], dtype=np.float32)
341    image_size = np.array(resize_dim)
342    rotation = 0
343
344    trans = get_affine_transform(center, scale, rotation, image_size)
345    warped_image = cv2.warpAffine(
346        image_np,
347        trans,
348        (int(image_size[0]), int(image_size[1])),
349        flags=cv2.INTER_LINEAR,
350    )
351
352    to_tensor = transforms.ToTensor()
353    normalize = transforms.Normalize(normalize_mean, normalize_std)
354    input_tensor = normalize(to_tensor(warped_image)).unsqueeze(0)
355
356    return input_tensor, image_np, center, scale