| import cv2
|
| import sys
|
| import argparse
|
| import numpy as np
|
| from PIL import Image
|
| from pathlib import Path
|
| from matplotlib import pyplot as plt
|
| from typing import Any, Dict, List
|
| import glob
|
|
|
| from utils import load_img_to_array, show_mask
|
|
|
|
|
| def setup_args(parser):
|
| parser.add_argument(
|
| "--input_img", type=str, required=True,
|
| help="Path to a single input img",
|
| )
|
| parser.add_argument(
|
| "--input_mask_glob", type=str, required=True,
|
| help="Glob to input masks",
|
| )
|
| parser.add_argument(
|
| "--output_dir", type=str, required=True,
|
| help="Output path to the directory with results.",
|
| )
|
|
|
| if __name__ == "__main__":
|
| """Example usage:
|
| python visual_mask_on_img.py \
|
| --input_img FA_demo/FA1_dog.png \
|
| --input_mask_glob "results/FA1_dog/mask*.png" \
|
| --output_dir results
|
| """
|
| parser = argparse.ArgumentParser()
|
| setup_args(parser)
|
| args = parser.parse_args(sys.argv[1:])
|
|
|
| img = load_img_to_array(args.input_img)
|
| img_stem = Path(args.input_img).stem
|
|
|
| mask_ps = sorted(glob.glob(args.input_mask_glob))
|
|
|
| out_dir = Path(args.output_dir) / img_stem
|
| out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
| for mask_p in mask_ps:
|
| mask = load_img_to_array(mask_p)
|
| mask = mask.astype(np.uint8)
|
|
|
|
|
| img_mask_p = out_dir / f"with_{Path(mask_p).name}"
|
|
|
|
|
| dpi = plt.rcParams['figure.dpi']
|
| height, width = img.shape[:2]
|
| plt.figure(figsize=(width/dpi/0.77, height/dpi/0.77))
|
| plt.imshow(img)
|
| plt.axis('off')
|
| show_mask(plt.gca(), mask, random_color=False)
|
| plt.savefig(img_mask_p, bbox_inches='tight', pad_inches=0)
|
| plt.close()
|
|
|