garmentiq.landmark.detect

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