garmentiq.classification.predict

  1"""Single-image garment classification."""
  2import torch
  3import torch.nn as nn
  4import torch.nn.functional as F
  5from PIL import Image
  6from torchvision import transforms
  7from typing import Type, List, Union
  8from garmentiq.utils.device import resolve_device, empty_cache
  9
 10
 11def predict(
 12    model: Type[nn.Module],
 13    image_path: str,
 14    classes: List[str],
 15    resize_dim=(120, 184),
 16    normalize_mean=[0.8047, 0.7808, 0.7769],
 17    normalize_std=[0.2957, 0.3077, 0.3081],
 18    device: Union[str, torch.device] = "cpu",
 19    verbose=False,
 20):
 21    """
 22    Loads a trained PyTorch model and makes a prediction on a single image.
 23
 24    This function processes a single image from disk, applies resizing and normalization,
 25    feeds it through a loaded model, and returns the predicted class along with the
 26    class probabilities. The model is expected to output logits over a fixed number of classes.
 27
 28    The model and the preprocessed image tensor are placed on `device`, so the same value should
 29    be passed here as was used when loading the model. Any accelerator memory cached during
 30    inference is released before returning.
 31
 32    Args:
 33        model (Type[nn.Module]): The loaded PyTorch model instance ready for inference.
 34        image_path (str): Path to the input image file (.jpg, .jpeg, .png).
 35        classes (List[str]): List of class names corresponding to model outputs. Will be sorted internally.
 36        resize_dim (tuple[int, int]): Tuple indicating the dimensions to resize the image to. Default is (120, 184).
 37        normalize_mean (list[float]): List of mean values for normalization. Default is [0.8047, 0.7808, 0.7769].
 38        normalize_std (list[float]): List of standard deviation values for normalization. Default is [0.2957, 0.3077, 0.3081].
 39        device (Union[str, torch.device], optional): The device to run inference on, e.g. `"cpu"`,
 40                                                     `"cuda"`, `"cuda:0"`, or `"mps"`. Hardware
 41                                                     acceleration is opt-in; pass it explicitly to
 42                                                     use a GPU or Apple Silicon. Defaults to `"cpu"`.
 43        verbose (bool): If True, prints the predicted label and class probabilities.
 44
 45    Raises:
 46        ValueError: If the image file does not have a supported extension (.jpg, .jpeg, .png),
 47                    or if the requested `device` is invalid or unavailable on this machine.
 48        FileNotFoundError: If the model checkpoint file is not found or cannot be loaded.
 49
 50    Returns:
 51        tuple[str, List[float]]: A tuple containing:
 52            - predicted label (str): The class label with the highest predicted probability.
 53            - prob_list (List[float]): The list of class probabilities in the same order as the sorted class list.
 54    """
 55    device = resolve_device(device)
 56
 57    # Validate image extension
 58    if not any(
 59        image_path.lower().endswith(ext) for ext in [".jpg", ".jpeg", ".png", ".JPG"]
 60    ):
 61        raise ValueError("Image file must end with .jpg, .jpeg, .png, or .JPG")
 62
 63    # Sort the classes list to have a consistent order
 64    sorted_classes = sorted(classes)
 65
 66    # Define the preprocessing transformation.
 67    transform = transforms.Compose(
 68        [
 69            transforms.Resize(resize_dim),
 70            transforms.ToTensor(),
 71            transforms.Normalize(mean=normalize_mean, std=normalize_std),
 72        ]
 73    )
 74
 75    # Load and preprocess the image
 76    image = Image.open(image_path).convert("RGB")
 77    model = model.to(device)
 78    image_tensor = transform(image).unsqueeze(0).to(device)  # add batch dimension
 79
 80    # Forward pass
 81    with torch.no_grad():
 82        outputs = model(image_tensor)
 83        # Compute probabilities using softmax
 84        probabilities = (
 85            F.softmax(outputs, dim=1).cpu().numpy()[0]
 86        )  # shape: (num_classes,)
 87
 88    # Determine the predicted index and label
 89    pred_index = int(probabilities.argmax())
 90    pred_label = sorted_classes[pred_index]
 91
 92    # Optionally, you might want to return probabilities as a list of floats:
 93    prob_list = probabilities.tolist()
 94
 95    if verbose:
 96        print(f"Prediction: {pred_label}")
 97        print(f"Probabilities: {prob_list}")
 98
 99    del image_tensor, outputs
100    empty_cache(device)
101
102    return pred_label, prob_list