garmentiq.landmark.detect

  1"""Detecting the predefined landmarks of a garment."""
  2import json
  3import os
  4from typing import Type, Union
  5import torch
  6import requests
  7import numpy as np
  8from garmentiq.utils import validate_garment_class_dict
  9from garmentiq.utils.device import resolve_device, empty_cache
 10from garmentiq.landmark.utils import (
 11    find_instruction_landmark_index,
 12    fill_instruction_landmark_coordinate,
 13)
 14from garmentiq.landmark.detection.utils import (
 15    input_image_transform,
 16    get_final_preds,
 17    transform_preds,
 18)
 19
 20
 21def detect(
 22    class_name: str,
 23    class_dict: dict,
 24    image_path: Union[str, np.ndarray],
 25    model: Type[torch.nn.Module],
 26    scale_std: float = 200.0,
 27    resize_dim: list[int, int] = [288, 384],
 28    normalize_mean: list[float, float, float] = [0.485, 0.456, 0.406],
 29    normalize_std: list[float, float, float] = [0.229, 0.224, 0.225],
 30    device: Union[str, torch.device] = "cpu",
 31):
 32    """
 33    Detects predefined landmarks on a garment image using a specified model and class instructions.
 34
 35    This function validates the input class dictionary and class name, loads the appropriate
 36    instruction schema (from local file or URL), preprocesses the image, runs it through
 37    the landmark detection model, and then transforms the detected heatmap predictions
 38    into image coordinates. The detected coordinates are then filled into the instruction data.
 39
 40    The model and the preprocessed input tensor are placed on `device`, so the same value should
 41    be passed here as was used when loading the model. Any accelerator memory cached during
 42    inference is released before returning.
 43
 44    Args:
 45        class_name (str): The name of the garment class (e.g., "vest dress", "trousers").
 46        class_dict (dict): A dictionary mapping class names to their properties, including
 47                           `num_predefined_points`, `index_range`, and `instruction` file path.
 48        image_path (Union[str, np.ndarray]): The path to the image file or a NumPy array of the image.
 49        model (Type[torch.nn.Module]): The loaded PyTorch landmark detection model.
 50        scale_std (float, optional): Standard scale for image transformation during preprocessing. Defaults to 200.0.
 51        resize_dim (list[int, int], optional): Target dimensions [width, height] for the transformed image.
 52                                               Defaults to [288, 384].
 53        normalize_mean (list[float, float, float], optional): Mean values for image normalization (RGB channels).
 54                                                              Defaults to [0.485, 0.456, 0.406].
 55        normalize_std (list[float, float, float], optional): Standard deviation values for image normalization (RGB channels).
 56                                                             Defaults to [0.229, 0.224, 0.225].
 57        device (Union[str, torch.device], optional): The device to run inference on, e.g. `"cpu"`,
 58                                                     `"cuda"`, `"cuda:0"`, or `"mps"`. Hardware
 59                                                     acceleration is opt-in; pass it explicitly to
 60                                                     use a GPU or Apple Silicon. Defaults to `"cpu"`.
 61
 62    Raises:
 63        ValueError: If `class_dict` is invalid, `class_name` is not found in `class_dict`, or the
 64                    requested `device` is invalid or unavailable on this machine.
 65        FileNotFoundError: If the instruction file is not found.
 66        ValueError: If loading instruction JSON from URL fails or `class_name` is not found in instruction file.
 67
 68    Returns:
 69        tuple:
 70            - preds_all (np.array): All predicted landmark coordinates (including non-predefined).
 71            - maxvals (np.array): Confidence scores for the predefined landmark predictions.
 72            - instruction_data (dict): The instruction dictionary updated with detected landmark coordinates and confidences.
 73    """
 74    device = resolve_device(device)
 75
 76    if not validate_garment_class_dict(class_dict):
 77        raise ValueError(
 78            "Provided class_dict is not in the expected garment_classes format."
 79        )
 80
 81    if class_name not in class_dict:
 82        raise ValueError(
 83            f"Invalid class '{class_name}'. Must be one of: {list(class_dict.keys())}"
 84        )
 85
 86    class_element = class_dict[class_name]
 87
 88    instruction_path = class_element["instruction"]
 89
 90    if instruction_path.startswith("http://") or instruction_path.startswith(
 91        "https://"
 92    ):
 93        try:
 94            response = requests.get(instruction_path)
 95            response.raise_for_status()
 96            instruction_data = response.json()
 97        except Exception as e:
 98            raise ValueError(
 99                f"Failed to load instruction JSON from URL: {instruction_path}\nError: {e}"
100            )
101    else:
102        if not os.path.exists(instruction_path):
103            raise FileNotFoundError(f"Instruction file not found: {instruction_path}")
104        with open(instruction_path, "r") as f:
105            instruction_data = json.load(f)
106
107    if class_name not in instruction_data:
108        raise ValueError(f"Class '{class_name}' not found in instruction file.")
109
110    (input_tensor, image_np, center, scale,) = input_image_transform(
111        image_path, scale_std, resize_dim, normalize_mean, normalize_std
112    )
113
114    model = model.to(device)
115    input_tensor = input_tensor.to(device)
116
117    with torch.no_grad():
118        np_output_heatmap = model(input_tensor).detach().cpu().numpy()
119
120    empty_cache(device)
121
122    preds_heatmap, maxvals = get_final_preds(
123        np_output_heatmap[
124            :, class_element["index_range"][0] : class_element["index_range"][1], :, :
125        ]
126    )
127
128    predefined_index = find_instruction_landmark_index(
129        instruction_data[class_name]["landmarks"], predefined=True
130    )
131    preds_all = np.stack([transform_preds(p, center, scale) for p in preds_heatmap])
132    preds = preds_all[:, predefined_index, :]
133
134    instruction_data[class_name]["landmarks"] = fill_instruction_landmark_coordinate(
135        instruction_landmarks=instruction_data[class_name]["landmarks"],
136        index=predefined_index,
137        fill_in_value=preds,
138    )
139
140    for idx in predefined_index:
141        instruction_data[class_name]["landmarks"][str(idx + 1)]["conf"] = float(
142            maxvals[0, idx, 0]
143        )
144
145    return preds_all, maxvals, instruction_data