|
from glob import glob |
|
import argparse |
|
import os |
|
from typing import Tuple, List |
|
from PIL import Image |
|
from rich.progress import track |
|
from vegseg.datasets import GrassDataset |
|
|
|
|
|
def get_args() -> Tuple[str, str, int]: |
|
""" |
|
get args |
|
return: |
|
--dataset_path: dataset path. |
|
--output_path: output path for saving. |
|
--num: num of image to show. -1 means all. |
|
""" |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--dataset_path", type=str, default="data/grass") |
|
parser.add_argument("--output_path", type=str, default="all_dataset.png") |
|
parser.add_argument("--num", default=-1, type=int, help="num of image to show") |
|
args = parser.parse_args() |
|
return args.dataset_path, args.output_path, args.num |
|
|
|
|
|
def get_image_and_mask_paths( |
|
dataset_path: str, num: int |
|
) -> Tuple[List[str], List[str]]: |
|
""" |
|
get image and mask paths from dataset path. |
|
return: |
|
image_paths: list of image paths. |
|
mask_paths: list of mask paths. |
|
""" |
|
image_paths = glob(os.path.join(dataset_path, "img_dir", "*", "*.tif")) |
|
if num != -1: |
|
image_paths = image_paths[:num] |
|
mask_paths = [ |
|
filename.replace("tif", "png").replace("img_dir", "ann_dir") |
|
for filename in image_paths |
|
] |
|
return image_paths, mask_paths |
|
|
|
|
|
def get_palette() -> List[int]: |
|
""" |
|
get palette of dataset. |
|
return: |
|
palette: list of palette. |
|
""" |
|
palette = [] |
|
palette_list = GrassDataset.METAINFO["palette"] |
|
for palette_item in palette_list: |
|
palette.extend(palette_item) |
|
return palette |
|
|
|
|
|
def paste_image_mask(image_path: str, mask_path: str) -> Image.Image: |
|
""" |
|
paste image and mask together |
|
Args: |
|
image_path (str): path to image. |
|
mask_path (str): path to mask. |
|
return: |
|
image_mask: image with mask,is Image. |
|
""" |
|
image = Image.open(image_path) |
|
mask = Image.open(mask_path).convert("P") |
|
palette = get_palette() |
|
mask.putpalette(palette) |
|
mask = mask.convert("RGB") |
|
image_mask = Image.new("RGB", (image.width * 2, image.height)) |
|
image_mask.paste(image, (0, 0)) |
|
image_mask.paste(mask, (image.width, 0)) |
|
return image_mask |
|
|
|
|
|
def paste_all_images(all_images: List[Image.Image], output_path: str) -> None: |
|
""" |
|
paste all images together and save it. |
|
Args: |
|
all_images (List[Image.Image]): list of image. |
|
output_path (str): path to save. |
|
Return: |
|
None |
|
""" |
|
widths = [image.width for image in all_images] |
|
heights = [image.height for image in all_images] |
|
width = max(widths) |
|
height = sum(heights) |
|
all_image = Image.new("RGB", (width, height)) |
|
for i, image in enumerate(all_images): |
|
all_image.paste(image, (0, sum(heights[:i]))) |
|
all_image.save(output_path) |
|
|
|
|
|
def main(): |
|
dataset_path, output_path, num = get_args() |
|
image_paths, mask_paths = get_image_and_mask_paths(dataset_path, num) |
|
all_images = [] |
|
for image_path, mask_path in zip(image_paths, mask_paths): |
|
image_mask = paste_image_mask(image_path, mask_path) |
|
all_images.append(image_mask) |
|
paste_all_images(all_images, output_path) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
main() |
|
|
