feat: suggest num repeat

Akegarasu · Sep 1, 2023 · 5b1da88 · 5b1da88
1 parent 5f0703e
commit 5b1da88
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/mikazuki/utils.py b/mikazuki/utils.py
@@ -33,7 +33,8 @@ def validate_data_dir(path):
         imgs = get_total_images(path, False)
         log.info(f"{len(imgs)} images found")
         if len(imgs) > 0:
-            dataset_path = os.path.join(path, "1_zkz")
+            num_repeat = suggest_num_repeat(len(imgs))
+            dataset_path = os.path.join(path, f"{num_repeat}_zkz")
             os.makedirs(dataset_path)
             for i in imgs:
                 shutil.move(i, dataset_path)
@@ -45,6 +46,17 @@ def validate_data_dir(path):
     return True
 
 
+def suggest_num_repeat(img_count):
+    if img_count < 10:
+        return 7
+    elif 10 < img_count <= 50:
+        return 5
+    elif 50 < img_count <= 100:
+        return 3
+
+    return 1
+
+
 def check_training_params(data):
     potential_path = [
         "train_data_dir", "reg_data_dir", "output_dir"