garmentiq.landmark.detection.utils

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