|
import os |
|
import click |
|
import json |
|
from tqdm import tqdm |
|
|
|
|
|
@click.command() |
|
@click.option("--swim_dir", type=str, default="datasets/swim_data") |
|
@click.option("--output_dir", type=str, default="datasets/swim_data_cyclegan") |
|
@click.option("--type", type=str, help="fog|rain|snow|night", required=True) |
|
@click.option("--no_night", is_flag=True) |
|
def build_cyclegan_dataset(swim_dir: str, output_dir: str, type: str, no_night: bool): |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
os.makedirs(output_dir, exist_ok=True) |
|
os.makedirs(os.path.join(output_dir, "trainA"), exist_ok=True) |
|
os.makedirs(os.path.join(output_dir, "trainB"), exist_ok=True) |
|
os.makedirs(os.path.join(output_dir, "testA"), exist_ok=True) |
|
os.makedirs(os.path.join(output_dir, "testB"), exist_ok=True) |
|
|
|
with open(os.path.join(swim_dir, "train", "labels.json"), "r") as f: |
|
train_labels = json.load(f) |
|
|
|
with open(os.path.join(swim_dir, "val", "labels.json"), "r") as f: |
|
val_labels = json.load(f) |
|
|
|
if type != "night": |
|
for label in tqdm(train_labels, desc="train"): |
|
if no_night and label["timeofday"] == "night": |
|
continue |
|
|
|
if label["weather"] == type: |
|
os.system( |
|
f"cp {os.path.join(swim_dir, 'train', 'images', label['name'])} {os.path.join(output_dir, 'trainB', label['name'])}" |
|
) |
|
elif label["weather"] == "clear": |
|
os.system( |
|
f"cp {os.path.join(swim_dir, 'train', 'images', label['name'])} {os.path.join(output_dir, 'trainA', label['name'])}" |
|
) |
|
|
|
for label in tqdm(val_labels, desc="val"): |
|
if no_night and label["timeofday"] == "night": |
|
continue |
|
|
|
if label["weather"] == type: |
|
os.system( |
|
f"cp {os.path.join(swim_dir, 'val', 'images', label['name'])} {os.path.join(output_dir, 'testB', label['name'])}" |
|
) |
|
elif label["weather"] == "clear": |
|
os.system( |
|
f"cp {os.path.join(swim_dir, 'val', 'images', label['name'])} {os.path.join(output_dir, 'testA', label['name'])}" |
|
) |
|
else: |
|
for label in tqdm(train_labels, desc="train"): |
|
if label["weather"] != "clear": |
|
continue |
|
|
|
if label["timeofday"] == "night": |
|
os.system( |
|
f"cp {os.path.join(swim_dir, 'train', 'images', label['name'])} {os.path.join(output_dir, 'trainB', label['name'])}" |
|
) |
|
elif label["timeofday"] == "daytime": |
|
os.system( |
|
f"cp {os.path.join(swim_dir, 'train', 'images', label['name'])} {os.path.join(output_dir, 'trainA', label['name'])}" |
|
) |
|
|
|
for label in tqdm(val_labels, desc="val"): |
|
if label["weather"] != "clear": |
|
continue |
|
|
|
if label["timeofday"] == "night": |
|
os.system( |
|
f"cp {os.path.join(swim_dir, 'val', 'images', label['name'])} {os.path.join(output_dir, 'testB', label['name'])}" |
|
) |
|
elif label["timeofday"] == "daytime": |
|
os.system( |
|
f"cp {os.path.join(swim_dir, 'val', 'images', label['name'])} {os.path.join(output_dir, 'testA', label['name'])}" |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
build_cyclegan_dataset() |
|
|