garmentiq.segmentation.process_and_save_images

  1"""Batch segmentation of a directory of images."""
  2import os
  3import torch
  4from transformers import AutoModelForImageSegmentation
  5from tqdm.auto import tqdm
  6from PIL import Image
  7import numpy as np
  8from typing import Union
  9from garmentiq.segmentation.extract import extract
 10from garmentiq.segmentation.change_background_color import change_background_color
 11
 12
 13def process_and_save_images(
 14    image_dir: str,
 15    output_dir: str,
 16    model: AutoModelForImageSegmentation,
 17    resize_dim: tuple[int, int],
 18    normalize_mean: list[float, float, float],
 19    normalize_std: list[float, float, float],
 20    background_color: tuple[int, int, int] = None,
 21    high_precision: bool = True,
 22    device: Union[str, torch.device] = "cpu",
 23):
 24    """
 25    Processes images from a directory by extracting segmentation masks and optionally modifying
 26    the background color, then saving the masks and modified images to specified output directories.
 27
 28    This function applies the `extract` function to segment each image, generating a mask, and
 29    optionally modifies the background of each image using the `change_background_color` function
 30    before saving both the masks and the modified images to disk.
 31
 32    Args:
 33        image_dir (str): The directory containing the input images to process.
 34        output_dir (str): The directory where the processed masks and modified images will be saved.
 35        model (transformers.AutoModelForImageSegmentation): The pre-trained model used for image segmentation.
 36        resize_dim (tuple[int, int]): The target dimensions to resize the images before processing (width, height).
 37        normalize_mean (list[float, float, float]): The mean values used for image normalization.
 38        normalize_std (list[float, float, float]): The standard deviation values used for image normalization.
 39        background_color (tuple[int, int, int], optional): The background color to apply to the image (RGB tuple),
 40                                                          or None to skip the modification. Default is None.
 41        high_precision (bool, optional): Whether to use high precision (32-bit) for image processing. Default is True.
 42        device (Union[str, torch.device], optional): The device to run inference on for every image, e.g.
 43                                                     `"cpu"`, `"cuda"`, `"cuda:0"`, or `"mps"`. Hardware
 44                                                     acceleration is opt-in; pass it explicitly to use a
 45                                                     GPU or Apple Silicon. Default is `"cpu"`.
 46
 47    Raises:
 48        FileNotFoundError: If the input image directory does not exist.
 49        ValueError: If the requested `device` is invalid or unavailable on this machine, or if the
 50                    `model` provided does not work correctly for image segmentation.
 51
 52    Returns:
 53        None. The processed masks and modified images are saved to the specified output directory.
 54    """
 55    # Create output directories for masks and modified images
 56    mask_dir = os.path.join(output_dir, "masks")
 57    os.makedirs(mask_dir, exist_ok=True)
 58
 59    if background_color is not None:
 60        modified_image_dir = os.path.join(output_dir, "bg_modified")
 61        os.makedirs(modified_image_dir, exist_ok=True)
 62
 63    image_files = [
 64        f
 65        for f in os.listdir(image_dir)
 66        if f.lower().endswith((".png", ".jpg", ".jpeg"))
 67    ]
 68
 69    # Loop through all image files in the provided directory
 70    with tqdm(total=len(image_files), desc="Processing Images", unit="image") as pbar:
 71        for filename in os.listdir(image_dir):
 72            if filename.lower().endswith((".png", ".jpg", ".jpeg")):
 73                image_path = os.path.join(image_dir, filename)
 74
 75                # Extract image and mask using the extract function
 76                image_np, mask_np = extract(
 77                    model=model,
 78                    image_path=image_path,
 79                    resize_dim=resize_dim,
 80                    normalize_mean=normalize_mean,
 81                    normalize_std=normalize_std,
 82                    high_precision=high_precision,
 83                    device=device,
 84                )
 85
 86                # Save the mask image
 87                mask_pil = Image.fromarray(mask_np)
 88                mask_pil.save(
 89                    os.path.join(mask_dir, f"mask_{os.path.splitext(filename)[0]}.png")
 90                )
 91
 92                if background_color is not None:
 93                    # Change the background color if provided
 94                    modified_image = change_background_color(
 95                        image_np, mask_np, background_color
 96                    )
 97
 98                    # Save the modified image with the new background color
 99                    modified_image_pil = Image.fromarray(modified_image)
100                    modified_image_pil.save(
101                        os.path.join(
102                            modified_image_dir,
103                            f"bg_modified_{os.path.splitext(filename)[0]}.png",
104                        )
105                    )
106
107                pbar.update(1)