Add files using upload-large-folder tool

frozenlake/data_process.py

@@ -234,23 +234,37 @@ def generate_dataset(
     with open(out / "path.json", "w") as f:
         json.dump(dict(sorted((s["image"], s["path_udrl"]) for s in all_samples)), f, indent=4)
 
+    # Stratified split: ensure each size is proportionally represented in test set
     random.seed(seed + 1)
-    random.shuffle(all_samples)
-    split = int(len(all_samples) * train_ratio)
-
-    def _jsonl(samples, path):
+    by_size: Dict[int, List[Dict]] = {}
+    for s in all_samples:
+        by_size.setdefault(s["maze_size"], []).append(s)
+
+    train_samples, test_samples = [], []
+    for sz in sorted(by_size):
+        group = by_size[sz]
+        random.shuffle(group)
+        sz_split = int(len(group) * train_ratio)
+        train_samples.extend(group[:sz_split])
+        test_samples.extend(group[sz_split:])
+
+    random.shuffle(train_samples)
+    random.shuffle(test_samples)
+    split = len(train_samples)
+
+    def _write_jsonl(samples, path):
         with open(path, "w") as f:
             for s in samples:
                 f.write(json.dumps(s) + "\n")
 
-    _jsonl(all_samples[:split], out / "train.jsonl")
-    _jsonl(all_samples[split:], out / "test.jsonl")
+    _write_jsonl(train_samples, out / "train.jsonl")
+    _write_jsonl(test_samples, out / "test.jsonl")
 
-    for name, samps in [("train", all_samples[:split]), ("test", all_samples[split:])]:
-        with open(out / f"{name}.csv", "w", newline="") as f:
+    for name, samples in [("train", train_samples), ("test", test_samples)]:
+        with open(out / f"{name}.csv", "w", newline="", encoding="utf-8") as f:
             w = csv.writer(f)
             w.writerow(["input_image", "video", "prompt"])
-            for s in samps:
+            for s in samples:
                 w.writerow([f"images/{s['image']}", f"videos/{s['video']}", s["prompt"]])
 
     state.size_progress = progress
@@ -382,8 +396,8 @@ def parse_args():
 
     gen = sub.add_parser("generate")
     gen.add_argument("--output-dir", default="frozenlake")
-    gen.add_argument("--sizes", type=int, nargs="+", default=[8, 16, 32])
-    gen.add_argument("--num-per-size", type=int, nargs="+", default=[100, 500, 1000])
+    gen.add_argument("--sizes", type=int, nargs="+", default=[8, 12, 16, 32])
+    gen.add_argument("--num-per-size", type=int, nargs="+", default=[1000, 2000, 5000, 10000])
     gen.add_argument("--p", type=float, default=0.5)
     gen.add_argument("--min-path-ratio", type=float, default=0.1)
     gen.add_argument("--img-size", type=int, default=1024)

maze/data_process.py

@@ -308,19 +308,33 @@ def generate_dataset(
     with open(out / "path.json", "w") as f:
         json.dump(dict(sorted(path_answers.items())), f, indent=4)
 
+    # Stratified split: ensure each size is proportionally represented in test set
     random.seed(seed + 1)
-    random.shuffle(all_samples)
-    split = int(len(all_samples) * train_ratio)
+    by_size: Dict[int, List[Dict]] = {}
+    for s in all_samples:
+        by_size.setdefault(s["maze_size"], []).append(s)
+
+    train_samples, test_samples = [], []
+    for sz in sorted(by_size):
+        group = by_size[sz]
+        random.shuffle(group)
+        sz_split = int(len(group) * train_ratio)
+        train_samples.extend(group[:sz_split])
+        test_samples.extend(group[sz_split:])
+
+    random.shuffle(train_samples)
+    random.shuffle(test_samples)
+    split = len(train_samples)
 
     def _write_jsonl(samples, path):
         with open(path, "w") as f:
             for s in samples:
                 f.write(json.dumps(s) + "\n")
 
-    _write_jsonl(all_samples[:split], out / "train.jsonl")
-    _write_jsonl(all_samples[split:], out / "test.jsonl")
+    _write_jsonl(train_samples, out / "train.jsonl")
+    _write_jsonl(test_samples, out / "test.jsonl")
 
-    for name, samples in [("train", all_samples[:split]), ("test", all_samples[split:])]:
+    for name, samples in [("train", train_samples), ("test", test_samples)]:
         with open(out / f"{name}.csv", "w", newline="", encoding="utf-8") as f:
             writer = csv.writer(f)
             writer.writerow(["input_image", "video", "prompt"])
@@ -354,6 +368,7 @@ def eval_videos(
     text_dir: str,
     output_json: Optional[str] = None,
     gt_json: Optional[str] = None,
+    strict: bool = True,
 ):
     """
     Evaluate a directory of result videos against ground-truth mazes.
@@ -401,6 +416,7 @@ def eval_videos(
 
     for vpath in tqdm(videos, desc="Extracting paths"):
         stem = vpath.stem                        # e.g. "size8_000"
+        stem = stem.replace('_gen', '') # Remove `_gen` suffix
         txt_path = txt_root / f"{stem}.txt"
 
         if not txt_path.exists():
@@ -442,7 +458,7 @@ def eval_videos(
         if maze is None:
             continue
         total_valid += 1
-        if proc.verify_path(maze["grid"], maze["start"], maze["end"], udrl):
+        if proc.verify_path(maze["grid"], maze["start"], maze["end"], udrl, strict=strict):
             correct += 1
             correctly_solved.append({"name": name, "length": len(udrl)})
 
@@ -471,7 +487,7 @@ def eval_videos(
         if sz not in size_stats:
             size_stats[sz] = {"total": 0, "correct": 0}
         size_stats[sz]["total"] += 1
-        if proc.verify_path(maze["grid"], maze["start"], maze["end"], udrl):
+        if proc.verify_path(maze["grid"], maze["start"], maze["end"], udrl, strict=strict):
             size_stats[sz]["correct"] += 1
 
     if size_stats:
@@ -490,7 +506,7 @@ def eval_videos(
 
     # Optional: compare with ground-truth JSON for path-length-binned accuracy
     if gt_json:
-        _compare_with_gt(extracted, gt_json, txt_root, proc)
+        _compare_with_gt(extracted, gt_json, txt_root, proc, strict=strict)
 
     print(f"{'=' * 50}")
 
@@ -500,6 +516,7 @@ def _compare_with_gt(
     gt_json_path: str,
     txt_root: Path,
     proc: MazeProcessor,
+    strict: bool = True,
 ):
     """Print accuracy binned by ground-truth path length."""
     try:
@@ -526,7 +543,7 @@ def _compare_with_gt(
 
         stem = name.replace(".png", "")
         maze = proc.load_text(str(txt_root / f"{stem}.txt"))
-        if maze and proc.verify_path(maze["grid"], maze["start"], maze["end"], pred_udrl):
+        if maze and proc.verify_path(maze["grid"], maze["start"], maze["end"], pred_udrl, strict=strict):
             bins[label]["correct"] += 1
 
     if bins:
@@ -539,7 +556,7 @@ def _compare_with_gt(
 
 # ==================== Verify: Pre-extracted JSON ====================
 
-def verify_results(json_file: str, text_dir: str):
+def verify_results(json_file: str, text_dir: str, strict: bool = True):
     """
     Verify pre-extracted UDRL paths (from a JSON file) against maze .txt files.
 
@@ -564,7 +581,7 @@ def verify_results(json_file: str, text_dir: str):
             skipped += 1
             continue
         valid += 1
-        if proc.verify_path(maze["grid"], maze["start"], maze["end"], udrl):
+        if proc.verify_path(maze["grid"], maze["start"], maze["end"], udrl, strict=strict):
             correct += 1
 
     acc = (correct / valid * 100) if valid else 0
@@ -586,9 +603,9 @@ def parse_args():
     # --- generate ---
     gen = sub.add_parser("generate", help="Generate dataset")
     gen.add_argument("--output-dir", type=str, default="maze")
-    gen.add_argument("--sizes", type=int, nargs="+", default=[8, 16, 24, 32])
-    gen.add_argument("--num-per-size", type=int, nargs="+", default=[100, 500, 1000, 2000])
-    gen.add_argument("--min-path-ratio", type=float, default=0.3,
+    gen.add_argument("--sizes", type=int, nargs="+", default=[8, 12, 16, 32])
+    gen.add_argument("--num-per-size", type=int, nargs="+", default=[1000, 1000, 1000, 2000])
+    gen.add_argument("--min-path-ratio", type=float, default=0.1,
                      help="Min path length as fraction of size²")
     gen.add_argument("--img-size", type=int, default=1024)
     gen.add_argument("--prompt", type=str,
@@ -616,13 +633,16 @@ def parse_args():
                     help="Output JSON for extracted paths (default: video_dir/0_result.json)")
     ev.add_argument("--gt-json", type=str, default=None,
                     help="Optional ground-truth path.json for length-binned accuracy")
+    ev.add_argument("--strict", action="store_true",
+                    help="Strict verification (exact UDRL match) vs leniency on no-op moves")
 
     # --- verify ---
     ver = sub.add_parser("verify", help="Verify a pre-extracted JSON of UDRL paths")
     ver.add_argument("json_file", type=str)
     ver.add_argument("--text-dir", type=str, required=True,
                      help="Directory with maze .txt files")
-
+    ver.add_argument("--strict", action="store_true",
+                     help="Strict verification (exact UDRL match) vs leniency on no-op moves")
     return p.parse_args()
 
 
@@ -639,10 +659,11 @@ if __name__ == "__main__":
             text_dir=args.text_dir,
             output_json=args.output_json,
             gt_json=args.gt_json,
+            strict=args.strict,
         )
 
     elif args.command == "verify":
-        verify_results(args.json_file, args.text_dir)
+        verify_results(args.json_file, args.text_dir, strict=args.strict)
 
     else:
         print("Usage: python maze_video_gen.py {generate|eval|verify} [options]")

maze/maze/checkpoints/Wan2.1-I2V-14B-720P_full_0211/epoch-0.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a8e11f69131acdcfd0a4914e924918814339f25179d5b3283813842028e95b30
+size 613510352

maze/maze/checkpoints/Wan2.1-I2V-14B-720P_lora_0216/epoch-0.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:48adc620a15f2093eaa2521cbd63c63ed8a5f7cb75a26b7ff7e723cbc91f0002
+size 1226928552

maze/maze/checkpoints/Wan2.1-I2V-14B-720P_lora_0216/epoch-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f561456a695d513b33524ff7b9a5a38754fa51b986402a4b10045d02f0b93aef
+size 1226928552

maze/maze/checkpoints/Wan2.1-I2V-14B-720P_lora_0216/epoch-2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2ef3d06a911780c81669fe80a9a5f756708263f877bddb3e9fb958b181b04c7
+size 1226928552

maze/maze_processor.py

@@ -162,7 +162,7 @@ class MazeProcessor:
 
     # ==================== Verification ====================
 
-    def verify_path(self, grid: Grid, start: Tuple, end: Tuple, udrl: str) -> bool:
+    def verify_path(self, grid: Grid, start: Tuple, end: Tuple, udrl: str, strict : bool = True) -> bool:
         """Verify that *udrl* is a wall-respecting walk from *start* to *end*."""
         n = len(grid)
         r, c = start
@@ -176,6 +176,9 @@ class MazeProcessor:
             if not (0 <= nr < n and 0 <= nc < n):
                 return False
             r, c = nr, nc
+            if not strict and (r, c) == end:
+                break
+
         return (r, c) == end
 
     # ==================== Text Encoding ====================

sudoku/generate_dataset.py

@@ -437,10 +437,23 @@ def generate_dataset(
                 f"{sum(1 for s in all_samples if s['clue'] == clue)} videos"
             )
 
-    # Final output
+    # Stratified split: ensure each clue level is proportionally represented
     random.seed(seed + 1)
-    random.shuffle(all_samples)
-    split_idx = int(len(all_samples) * train_ratio)
+    by_clue: Dict[int, List[Dict]] = {}
+    for s in all_samples:
+        by_clue.setdefault(s["clue"], []).append(s)
+
+    train_samples, test_samples = [], []
+    for clue in sorted(by_clue):
+        group = by_clue[clue]
+        random.shuffle(group)
+        cl_split = int(len(group) * train_ratio)
+        train_samples.extend(group[:cl_split])
+        test_samples.extend(group[cl_split:])
+
+    random.shuffle(train_samples)
+    random.shuffle(test_samples)
+    split_idx = len(train_samples)
 
     def write_jsonl(samples, path):
         with open(path, 'w') as f:
@@ -448,8 +461,8 @@ def generate_dataset(
                 json.dump(s, f)
                 f.write('\n')
 
-    write_jsonl(all_samples[:split_idx], output_dir / "train.jsonl")
-    write_jsonl(all_samples[split_idx:], output_dir / "test.jsonl")
+    write_jsonl(train_samples, output_dir / "train.jsonl")
+    write_jsonl(test_samples, output_dir / "test.jsonl")
 
     # Mark as completed
     state.clue_progress = clue_progress