cc

.gitignore

@@ -1,2 +1,3 @@
 __pycache__
-datasets
+datasets
+*.zip

scripts/build_cyclegan_dataset.py

@@ -0,0 +1,67 @@
+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)
+def build_cyclegan_dataset(swim_dir: str, output_dir: str, type: str):
+    # 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)
+
+    for label in tqdm(train_labels, desc="train"):
+        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 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'])}"
+            )
+
+
+if __name__ == "__main__":
+    build_cyclegan_dataset()

scripts/build_swim_dataset.py

@@ -162,8 +162,8 @@ def build_swim_dataset(bdd100k_dir: str, acdc_dir: str, output_dir: str):
                         timeofday_ = "daytime"
                         weather_ = "clear"
                     else:
-                        timeofday_ = "daytime" if weather == "night" else "night"
-                        weather_ = weather
+                        timeofday_ = "night" if weather == "night" else "daytime"
+                        weather_ = "clear" if weather == "night" else weather
 
                     labels.append(
                         {
@@ -228,8 +228,8 @@ def build_swim_dataset(bdd100k_dir: str, acdc_dir: str, output_dir: str):
                         timeofday_ = "daytime"
                         weather_ = "clear"
                     else:
-                        timeofday_ = "daytime" if weather == "night" else "night"
-                        weather_ = weather
+                        timeofday_ = "night" if weather == "night" else "daytime"
+                        weather_ = "clear" if weather == "night" else weather
 
                     labels.append(
                         {