Merge pull request #13 from script-money/fix/space

1. feat: check all images then show all problem images
2. fix: two parts with string contain relation will always select first part
3. feat: add weight and trait folder validator

Author: script-money

src/generate.py

@@ -6,6 +6,7 @@
 import random
 from numpy.random import choice
 from get_table import files_path
+from pathlib import Path
 from config import (
     W,
     H,
@@ -112,9 +113,9 @@ def generate_func(
             next(
                 path
                 for path in files_path
-                if attr["trait_type"] in path.split(os.sep)[1]
-                and attr["value"][0] in path.split(os.sep)[0]
-                and attr["value"][2] in path.split(os.sep)[2]
+                if attr["trait_type"] == path.split(os.sep)[1].split("_")[1]
+                and attr["value"][0] == path.split(os.sep)[0]
+                and attr["value"][2] == Path(path).stem
             )
             for attr in attributes
         ]

src/get_table.py

@@ -2,7 +2,7 @@
 from PIL import Image
 import pandas as pd
 import shutil
-from config import FOLDERS, EXTENSION, W, H
+from config import FOLDERS, EXTENSION, W, H, WEIGHTS
 
 
 def get_files_path(folders=FOLDERS):
@@ -44,22 +44,58 @@ def get_files_path(folders=FOLDERS):
             print(f"{folder} folder cleaned")
         os.mkdir(folder)
 
+    # Validate weights
+    assert (
+        sum(WEIGHTS) == 1
+    ), f"sum of PARTS_DICT's value in config.py should be 1, now is {sum(WEIGHTS)}"
+
     # Validate image format and size
+    error = 0
     for path in files_path:
-        assert (
-            path.split(".")[-1] in EXTENSION
-        ), f"{path}'s extension is not {EXTENSION} "
-        im = Image.open(path)
-        w, h = im.size
-        assert w == W, f"{path} width not equal {W}"
-        assert h == H, f"{path} height not equal {H}"
+        try:
+            assert (
+                path.split(".")[-1] in EXTENSION
+            ), f"{path}'s extension is not {EXTENSION} "
+            im = Image.open(path)
+            w, h = im.size
+            assert w == W, f"{path} width not equal {W}"
+            assert h == H, f"{path} height not equal {H}"
+        except Exception as e:
+            print(e)
+            error += 1
+    if error != 0:
+        print(f"{error} images have error, fix and try again")
+        exit()
 
     # Validate path name has -
     for path in files_path:
         assert (
             "-" not in path.split("/")[-1]
         ), f"{path} is invalid, files should not have '-' symbol"
 
+    folder_set = set()
+    folder_error = 0
+    # check all parts folders has same order
+    for folder in FOLDERS:
+        try:
+            for root, subfolders, _ in os.walk(folder):
+                if root == folder:
+                    for subfolder in subfolders:
+                        if subfolder not in folder_set and subfolder.split("_")[
+                            1
+                        ] in map(lambda f: f.split("_")[1], folder_set):
+                            raise Exception(
+                                f"in '{root}' folder '{subfolder}' is invalid, index of subfolder with same trait should be same."
+                            )
+                        else:
+                            folder_set.add(subfolder)
+        except Exception as e:
+            print(e)
+            folder_error += 1
+    if folder_error != 0:
+        print("Exited, please fix subfolder error and retry")
+        exit()
+
     # export tables
     attrs = [os.path.split(path) for path in files_path]
     d = {