garmentiq.landmark.refine

 1"""Refining detected landmarks against a segmentation mask."""
 2import numpy as np
 3import cv2
 4import copy
 5from garmentiq.landmark.refinement import refine_landmark_with_blur
 6from garmentiq.landmark.utils import (
 7    find_instruction_landmark_index,
 8    fill_instruction_landmark_coordinate,
 9)
10
11
12def refine(
13    class_name: str,
14    detection_np: np.array,
15    detection_conf: np.array,
16    detection_dict: dict,
17    mask: np.array,
18    window_size: int = 5,
19    ksize: tuple = (11, 11),
20    sigmaX: float = 0.0,
21):
22    """
23    Refines detected landmarks using a blurred mask and updates the detection dictionary.
24
25    This function applies Gaussian blur to the given mask, then refines landmark coordinates
26    based on their confidence scores and local intensity structure. Only landmarks with a
27    confidence score greater than 0 are refined. The refined coordinates are used to update
28    predefined landmarks in the detection dictionary.
29
30    Args:
31        class_name (str): The name of the class to access in the detection dictionary.
32        detection_np (np.array): The initial landmark predictions. Shape: (1, N, 2).
33        detection_conf (np.array): Confidence scores for each predicted landmark. Shape: (1, N, 1).
34        detection_dict (dict): Dictionary containing landmark data for each class.
35        mask (np.array): Grayscale mask image used to guide refinement.
36        window_size (int, optional): Size of the window used in the refinement algorithm. Defaults to 5.
37        ksize (tuple, optional): Kernel size for Gaussian blur. Must be odd integers. Defaults to (11, 11).
38        sigmaX (float, optional): Gaussian kernel standard deviation in the X direction. Defaults to 0.0.
39
40    Returns:
41        tuple:
42            - refined_detection_np (np.array): Array of the same shape as `detection_np` with refined coordinates.
43            - detection_dict (dict): Updated detection dictionary with refined landmark coordinates.
44    """
45    blurred_mask = cv2.GaussianBlur(mask, ksize, sigmaX)
46
47    refined_detection_np = np.zeros(detection_np.shape)
48
49    for i, coord in enumerate(detection_np[0]):
50        if detection_conf[0, i, 0] > 0:
51            refined_x, refined_y = refine_landmark_with_blur(
52                coord[0], coord[1], blurred_mask, window_size
53            )
54            refined_detection_np[0, i] = [refined_x, refined_y]
55
56    predefined_index = find_instruction_landmark_index(
57        detection_dict[class_name]["landmarks"], predefined=True
58    )
59
60    preds = refined_detection_np[:, predefined_index, :]
61
62    detection_dict_copy = copy.deepcopy(detection_dict)
63
64    detection_dict_copy[class_name]["landmarks"] = fill_instruction_landmark_coordinate(
65        instruction_landmarks=detection_dict_copy[class_name]["landmarks"],
66        index=predefined_index,
67        fill_in_value=preds,
68    )
69
70    return refined_detection_np, detection_dict_copy