| |
| |
| |
| |
| |
| |
|
|
| import sys |
|
|
| sys.path.append(".") |
|
|
| import os |
| import pdb |
| import time |
| import traceback |
|
|
| import cv2 |
| import numpy as np |
| from tqdm import tqdm |
|
|
| from engine.SegmentAPI.SAM import SAM2Seg, erode_and_dialted, eroded, fill_mask |
|
|
|
|
| def basename(path): |
|
|
| pre_name = os.path.basename(path).split(".")[0] |
|
|
| return pre_name |
|
|
|
|
| def multi_process(worker, items, **kwargs): |
| """ |
| worker worker function to process items |
| """ |
|
|
| nodes = kwargs["nodes"] |
| dirs = kwargs["dirs"] |
|
|
| bucket = int( |
| np.ceil(len(items) / nodes) |
| ) |
|
|
| print("Total Nodes:", nodes) |
| print("Save Path:", dirs) |
| rank = int(os.environ.get("RANK", 0)) |
|
|
| print("Current Rank:", rank) |
|
|
| kwargs["RANK"] = rank |
|
|
| if rank == nodes - 1: |
| output_dir = worker(items[bucket * rank :], **kwargs) |
| else: |
| output_dir = worker(items[bucket * rank : bucket * (rank + 1)], **kwargs) |
|
|
| if rank == 0 and nodes > 1: |
| timesleep = int(kwargs.get("sleep", 7200)) |
| time.sleep(timesleep) |
|
|
|
|
| def is_img(image_path): |
| return os.path.isfile(image_path) and image_path.lower().endswith( |
| (".png", ".jpg", ".jpeg", ".gif") |
| ) |
|
|
|
|
| def run_seg(items, **params): |
| output_dir = params["dirs"] |
| debug = params["debug"] |
| model = params["model"] |
| os.makedirs(output_dir, exist_ok=True) |
|
|
| if debug: |
| items = items[:5] |
|
|
| process_valid = [] |
|
|
| for item in tqdm(items, desc="Processing..."): |
|
|
| if not os.path.isdir(item): |
| continue |
|
|
| print(f"processing img dir {item}") |
| basename_folder = basename(item) |
| files = os.listdir(item) |
| files = [os.path.join(item, file) for file in files] |
|
|
| img_files = list(filter(lambda x: is_img(x), files)) |
|
|
| save_folder = os.path.join(output_dir, basename_folder) |
| os.makedirs(os.path.join(output_dir, basename_folder), exist_ok=True) |
|
|
| try: |
| for img_file in img_files: |
| base_file_name = basename(img_file) |
|
|
| out = model(img_path=img_file, bbox=None) |
| alpha = fill_mask(out.masks) |
| alpha = erode_and_dialted( |
| (alpha * 255).astype(np.uint8), kernel_size=3, iterations=3 |
| ) |
| save_path = os.path.join(save_folder, base_file_name + ".jpg") |
| cv2.imwrite(save_path, alpha) |
| except: |
| continue |
|
|
| process_valid.append(item) |
|
|
| return output_dir |
|
|
|
|
| def get_parse(): |
| import argparse |
|
|
| parser = argparse.ArgumentParser(description="") |
| parser.add_argument("-i", "--input", required=True, help="input path") |
| parser.add_argument("-o", "--output", required=True, help="output path") |
| parser.add_argument("--nodes", default=1, type=int, help="how many workload?") |
| parser.add_argument("--debug", action="store_true", help="debug tag") |
| parser.add_argument("--txt", default=None, type=str) |
| args = parser.parse_args() |
| return args |
|
|
|
|
| if __name__ == "__main__": |
| opt = get_parse() |
|
|
| |
|
|
| if opt.txt == None: |
| available_items = os.listdir(opt.input) |
| available_items = [os.path.join(opt.input, item) for item in available_items] |
| else: |
| available_items = [] |
| with open(opt.txt) as reader: |
| for line in reader: |
| available_items.append(line.strip()) |
|
|
| available_items = [ |
| os.path.join(opt.input, item.split(".")[0]) for item in available_items |
| ] |
|
|
| prior_seg = SAM2Seg(wo_supres=False) |
|
|
| multi_process( |
| worker=run_seg, |
| items=available_items, |
| dirs=opt.output, |
| nodes=opt.nodes, |
| debug=opt.debug, |
| model=prior_seg, |
| ) |
|
|
