File size: 3,951 Bytes
4a0bcb8 9b66f69 4a0bcb8 9b66f69 be91bac 9b66f69 4a0bcb8 9b66f69 be91bac 9b66f69 798fdd3 be91bac 9b66f69 be91bac 9b66f69 798fdd3 be91bac 9b66f69 be91bac 9b66f69 4a0bcb8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
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):
# build the dataset with format
# swim_data_cyclegan
# ├── trainA
# │ ├── 000000.jpg
# │ ├── ...
# ├── trainB
# │ ├── 000000.jpg
# │ ├── ...
# |── testA
# │ ├── 000000.jpg
# │ ├── ...
# |── testB
# │ ├── 000000.jpg
# │ ├── ...
#
# note:
# trainA and trainB are extracted from swim_data train set
# testA and testB are extracted from swim_data val set
# A - clear
# B - fog|rain|snow|night
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()
|