garmentiq.segmentation.extract

 1from PIL import Image
 2import torch
 3from torchvision import transforms
 4import numpy as np
 5
 6
 7def extract(model: torch.nn.Module, image_path: str, processor=None, **kwargs):
 8    """
 9    Intelligently extracts an image segmentation mask from a given image using either a standard
10    PyTorch model or a Processor-based foundation model.
11
12    This function takes an image and processes it based on the model strategy. If a processor is supplied,
13    it delegates preprocessing (e.g., resizing, scaling, prompt-handling) to the processor. Otherwise,
14    it applies standard manual transformations based on provided kwargs. It then feeds the input into
15    the model to generate a segmentation mask. The original image and the mask are returned as numpy arrays.
16
17    Args:
18        model (torch.nn.Module): The pretrained PyTorch model to use for segmentation predictions.
19        image_path (str): The path to the image file on which to perform segmentation.
20        processor (Any, optional): The model-specific processor (e.g., from Hugging Face) used for preprocessing
21                                   inputs. If None, standard PyTorch manual transformations are applied.
22                                   Default is None.
23        **kwargs: Additional arbitrary keyword arguments for model-specific configurations.
24                  For standard models (e.g., BiRefNet): `resize_dim`, `normalize_mean`, `normalize_std`, `output_index`, `output_key`.
25                  For processor models (e.g., SAM): Specific prompt arguments like `input_points`.
26
27    Raises:
28        FileNotFoundError: If the image file at `image_path` does not exist.
29        ValueError: If the model is incompatible with the task or the processor output format is unrecognized.
30
31    Returns:
32        tuple (numpy.ndarray, numpy.ndarray): The original image converted to a numpy array,
33                                              and the extracted segmentation mask as a numpy array.
34    """
35    device = "cuda" if next(model.parameters()).is_cuda else "cpu"
36    image = Image.open(image_path).convert("RGB")
37
38    # Processor-Based Models (SAM)
39    if processor is not None:
40        # Pass the image and any SAM-specific kwargs (like input_points) to the processor
41        inputs = processor(image, return_tensors="pt", **kwargs).to(device)
42
43        with torch.no_grad():
44            outputs = model(**inputs)
45
46        # Handle SAM-specific output format
47        if hasattr(outputs, "pred_masks"):
48            masks = processor.image_processor.post_process_masks(
49                outputs.pred_masks.cpu(),
50                inputs["original_sizes"].cpu(),
51                inputs["reshaped_input_sizes"].cpu(),
52            )
53            # Extract the best mask for the first point
54            best_mask = masks[0][0][0].numpy()
55            mask_np = (best_mask * 255).astype(np.uint8)
56
57        else:
58            raise ValueError("Unrecognized processor output format.")
59
60        del inputs, outputs, masks
61
62    # Standard Models (BiRefNet)
63    else:
64        # Extract BiRefNet-specific kwargs with safe defaults
65        resize_dim = kwargs.get("resize_dim", (1024, 1024))
66        normalize_mean = kwargs.get("normalize_mean", [0.485, 0.456, 0.406])
67        normalize_std = kwargs.get("normalize_std", [0.229, 0.224, 0.225])
68
69        transform = transforms.Compose(
70            [
71                transforms.Resize(resize_dim),
72                transforms.ToTensor(),
73                transforms.Normalize(normalize_mean, normalize_std),
74            ]
75        )
76
77        input_tensor = transform(image).unsqueeze(0).to(device)
78
79        with torch.no_grad():
80            preds = model(input_tensor)
81
82            # BiRefNet returns a tuple/list of tensors; we want the last one
83            if isinstance(preds, (list, tuple)):
84                preds = preds[-1]
85
86            preds = preds.sigmoid().cpu()
87
88        pred = preds[0].squeeze()
89        pred_pil = transforms.ToPILImage()(pred)
90        mask = pred_pil.resize(image.size)
91        mask_np = np.array(mask)
92
93        del input_tensor, preds
94
95    # Clean up and Return
96    image_np = np.array(image)
97    torch.cuda.empty_cache()
98
99    return image_np, mask_np