garmentiq.segmentation.process_and_save_images

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