From 95beb8a64532a2c4a554dd714c672e5bfbb4c827 Mon Sep 17 00:00:00 2001 From: Federico Baldassarre Date: Thu, 25 May 2023 17:10:44 +0200 Subject: [PATCH] Improve scripts/amg.py: paths, progress bar, multi-GPU, resuming - Change all os.path methods to the nicer pathlib.Path - Respect subfolders when processing a directory e.g. my/dataset/class_{0,1,2}/img_{0,1,2}.jpg -> my/masks/class_{0,1,2}/img_{0,1,2}.json - Show progress bar using tqdm - Load images in a separate dataloader process - Split a large dataset across multiple GPUs using rank/world size - Skip processing images that already have a corresponding output dir (png masks) or json file (rle masks) --- scripts/amg.py | 99 +++++++++++++++++++++++++++++++++++--------------- 1 file changed, 69 insertions(+), 30 deletions(-) diff --git a/scripts/amg.py b/scripts/amg.py index f2dbf676a..93c96086a 100644 --- a/scripts/amg.py +++ b/scripts/amg.py @@ -4,15 +4,17 @@ # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. -import cv2 # type: ignore - -from segment_anything import SamAutomaticMaskGenerator, sam_model_registry - import argparse import json -import os +from pathlib import Path from typing import Any, Dict, List +import cv2 # type: ignore +import tqdm +import torch.utils.data + +from segment_anything import SamAutomaticMaskGenerator, sam_model_registry + parser = argparse.ArgumentParser( description=( "Runs automatic mask generation on an input image or directory of images, " @@ -53,6 +55,9 @@ ) parser.add_argument("--device", type=str, default="cuda", help="The device to run generation on.") +parser.add_argument("--rank", type=int, default=0, help="Rank of the current process.") +parser.add_argument("--world", type=int, default=1, help="Number of processes.") +parser.add_argument("--num-workers", type=int, default=4, help="Dataloader workers.") parser.add_argument( "--convert-to-rle", @@ -149,13 +154,13 @@ ) -def write_masks_to_folder(masks: List[Dict[str, Any]], path: str) -> None: +def write_masks_to_folder(masks: List[Dict[str, Any]], path: Path) -> None: header = "id,area,bbox_x0,bbox_y0,bbox_w,bbox_h,point_input_x,point_input_y,predicted_iou,stability_score,crop_box_x0,crop_box_y0,crop_box_w,crop_box_h" # noqa metadata = [header] for i, mask_data in enumerate(masks): mask = mask_data["segmentation"] filename = f"{i}.png" - cv2.imwrite(os.path.join(path, filename), mask * 255) + cv2.imwrite((path / filename).as_posix(), mask * 255) mask_metadata = [ str(i), str(mask_data["area"]), @@ -167,8 +172,7 @@ def write_masks_to_folder(masks: List[Dict[str, Any]], path: str) -> None: ] row = ",".join(mask_metadata) metadata.append(row) - metadata_path = os.path.join(path, "metadata.csv") - with open(metadata_path, "w") as f: + with open(path / "metadata.csv", "w") as f: f.write("\n".join(metadata)) return @@ -192,6 +196,21 @@ def get_amg_kwargs(args): return amg_kwargs +class ImageDataset(torch.utils.data.Dataset): + def __init__(self, paths: List[Path], base: Path): + self.paths = paths + self.base = base + + def __len__(self): + return len(self.paths) + + def __getitem__(self, idx): + path = self.paths[idx] + image = cv2.imread((self.base / path).as_posix()) + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + return path, image + + def main(args: argparse.Namespace) -> None: print("Loading model...") sam = sam_model_registry[args.model_type](checkpoint=args.checkpoint) @@ -200,34 +219,54 @@ def main(args: argparse.Namespace) -> None: amg_kwargs = get_amg_kwargs(args) generator = SamAutomaticMaskGenerator(sam, output_mode=output_mode, **amg_kwargs) - if not os.path.isdir(args.input): - targets = [args.input] + args.input = Path(args.input).expanduser().resolve() + args.output = Path(args.output).expanduser().resolve() + + if not args.input.is_dir(): + targets = ImageDataset([args.input.name], args.input.parent) else: targets = [ - f for f in os.listdir(args.input) if not os.path.isdir(os.path.join(args.input, f)) + f.relative_to(args.input) for f in args.input.rglob("*") + if f.is_file() and f.suffix.lower() in [".jpg", ".jpeg", ".png"] ] - targets = [os.path.join(args.input, f) for f in targets] - - os.makedirs(args.output, exist_ok=True) - - for t in targets: - print(f"Processing '{t}'...") - image = cv2.imread(t) - if image is None: - print(f"Could not load '{t}' as an image, skipping...") - continue - image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) - + print(f"Found {len(targets)} images in {args.input}.") + + # Per-process split + if args.world > 1: + targets = targets[args.rank::args.world] + print(f"Rank {args.rank}/{args.world} will process {len(targets)} images.") + + # Skip existing + if output_mode == "binary_mask": + targets = [ + f for f in targets + if not Path.is_dir(args.output / f.with_suffix("")) + ] + else: + targets = [ + f for f in targets + if not Path.is_file(args.output / f.with_suffix(".json")) + ] + print(f"Skip already processed images, {len(targets)} remain to do.") + + targets = torch.utils.data.DataLoader( + ImageDataset(targets, args.input), + batch_size=None, + shuffle=False, + num_workers=args.num_workers, + collate_fn=lambda x: x, + ) + + for path, image in tqdm.tqdm(targets, ncols=0): masks = generator.generate(image) - base = os.path.basename(t) - base = os.path.splitext(base)[0] - save_base = os.path.join(args.output, base) if output_mode == "binary_mask": - os.makedirs(save_base, exist_ok=False) - write_masks_to_folder(masks, save_base) + save_dir = args.output / path.with_suffix("") + save_dir.mkdir(parents=True, exist_ok=False) + write_masks_to_folder(masks, save_dir) else: - save_file = save_base + ".json" + save_file = args.output / path.with_suffix(".json") + save_file.parent.mkdir(parents=True, exist_ok=True) with open(save_file, "w") as f: json.dump(masks, f) print("Done!")