halib 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (59) hide show
  1. halib/__init__.py +3 -3
  2. halib/common/__init__.py +0 -0
  3. halib/common/common.py +178 -0
  4. halib/common/rich_color.py +285 -0
  5. halib/exp/__init__.py +0 -0
  6. halib/exp/core/__init__.py +0 -0
  7. halib/exp/core/base_config.py +144 -0
  8. halib/exp/core/base_exp.py +157 -0
  9. halib/exp/core/param_gen.py +108 -0
  10. halib/exp/core/wandb_op.py +117 -0
  11. halib/exp/data/__init__.py +0 -0
  12. halib/exp/data/dataclass_util.py +41 -0
  13. halib/exp/data/dataset.py +208 -0
  14. halib/exp/data/torchloader.py +165 -0
  15. halib/exp/perf/__init__.py +0 -0
  16. halib/exp/perf/flop_calc.py +190 -0
  17. halib/exp/perf/gpu_mon.py +58 -0
  18. halib/exp/perf/perfcalc.py +363 -0
  19. halib/exp/perf/perfmetrics.py +137 -0
  20. halib/exp/perf/perftb.py +778 -0
  21. halib/exp/perf/profiler.py +301 -0
  22. halib/exp/viz/__init__.py +0 -0
  23. halib/exp/viz/plot.py +754 -0
  24. halib/filetype/csvfile.py +3 -9
  25. halib/filetype/ipynb.py +3 -5
  26. halib/filetype/jsonfile.py +0 -3
  27. halib/filetype/textfile.py +0 -1
  28. halib/filetype/videofile.py +91 -2
  29. halib/filetype/yamlfile.py +3 -3
  30. halib/online/projectmake.py +7 -6
  31. halib/online/tele_noti.py +165 -0
  32. halib/research/core/__init__.py +0 -0
  33. halib/research/core/base_config.py +144 -0
  34. halib/research/core/base_exp.py +157 -0
  35. halib/research/core/param_gen.py +108 -0
  36. halib/research/core/wandb_op.py +117 -0
  37. halib/research/data/__init__.py +0 -0
  38. halib/research/data/dataclass_util.py +41 -0
  39. halib/research/data/dataset.py +208 -0
  40. halib/research/data/torchloader.py +165 -0
  41. halib/research/perf/__init__.py +0 -0
  42. halib/research/perf/flop_calc.py +190 -0
  43. halib/research/perf/gpu_mon.py +58 -0
  44. halib/research/perf/perfcalc.py +363 -0
  45. halib/research/perf/perfmetrics.py +137 -0
  46. halib/research/perf/perftb.py +778 -0
  47. halib/research/perf/profiler.py +301 -0
  48. halib/research/viz/__init__.py +0 -0
  49. halib/research/viz/plot.py +754 -0
  50. halib/system/filesys.py +60 -20
  51. halib/system/path.py +73 -0
  52. halib/utils/dict.py +9 -0
  53. halib/utils/list.py +12 -0
  54. {halib-0.2.1.dist-info → halib-0.2.3.dist-info}/METADATA +5 -2
  55. halib-0.2.3.dist-info/RECORD +108 -0
  56. halib-0.2.1.dist-info/RECORD +0 -64
  57. {halib-0.2.1.dist-info → halib-0.2.3.dist-info}/WHEEL +0 -0
  58. {halib-0.2.1.dist-info → halib-0.2.3.dist-info}/licenses/LICENSE.txt +0 -0
  59. {halib-0.2.1.dist-info → halib-0.2.3.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,157 @@
1
+ from abc import ABC, abstractmethod
2
+ from typing import Tuple, Any, Optional
3
+ from base_config import ExpBaseConfig
4
+ from ..perf.perfcalc import PerfCalc
5
+ from ..perf.perfmetrics import MetricsBackend
6
+
7
+ # ! SEE https://github.com/hahv/base_exp for sample usage
8
+ class BaseExperiment(PerfCalc, ABC):
9
+ """
10
+ Base class for experiments.
11
+ Orchestrates the experiment pipeline using a pluggable metrics backend.
12
+ """
13
+
14
+ def __init__(self, config: ExpBaseConfig):
15
+ self.config = config
16
+ self.metric_backend = None
17
+ # Flag to track if init_general/prepare_dataset has run
18
+ self._is_env_ready = False
19
+
20
+ # -----------------------
21
+ # PerfCalc Required Methods
22
+ # -----------------------
23
+ def get_dataset_name(self):
24
+ return self.config.get_dataset_cfg().get_name()
25
+
26
+ def get_experiment_name(self):
27
+ return self.config.get_cfg_name()
28
+
29
+ def get_metric_backend(self):
30
+ if not self.metric_backend:
31
+ self.metric_backend = self.prepare_metrics(self.config.get_metric_cfg())
32
+ return self.metric_backend
33
+
34
+ # -----------------------
35
+ # Abstract Experiment Steps
36
+ # -----------------------
37
+ @abstractmethod
38
+ def init_general(self, general_cfg):
39
+ """Setup general settings like SEED, logging, env variables."""
40
+ pass
41
+
42
+ @abstractmethod
43
+ def prepare_dataset(self, dataset_cfg):
44
+ """Load/prepare dataset."""
45
+ pass
46
+
47
+ @abstractmethod
48
+ def prepare_metrics(self, metric_cfg) -> MetricsBackend:
49
+ """
50
+ Prepare the metrics for the experiment.
51
+ This method should be implemented in subclasses.
52
+ """
53
+ pass
54
+
55
+ @abstractmethod
56
+ def before_exec_exp_once(self, *args, **kwargs):
57
+ """Optional: any setup before exec_exp. Note this is called once per run_exp."""
58
+ pass
59
+
60
+ @abstractmethod
61
+ def exec_exp(self, *args, **kwargs) -> Optional[Tuple[Any, Any]]:
62
+ """Run experiment process, e.g.: training/evaluation loop.
63
+ Return: either `None` or a tuple of (raw_metrics_data, extra_data) for calc_and_save_exp_perfs
64
+ """
65
+ pass
66
+
67
+ @abstractmethod
68
+ def exec_eval(self, *args, **kwargs) -> Optional[Tuple[Any, Any]]:
69
+ """Run evaluation process.
70
+ Return: either `None` or a tuple of (raw_metrics_data, extra_data) for calc_and_save_exp_perfs
71
+ """
72
+ pass
73
+
74
+ # -----------------------
75
+ # Internal Helpers
76
+ # -----------------------
77
+ def _validate_and_unpack(self, results):
78
+ if results is None:
79
+ return None
80
+ if not isinstance(results, (tuple, list)) or len(results) != 2:
81
+ raise ValueError("exec must return (metrics_data, extra_data)")
82
+ return results[0], results[1]
83
+
84
+ def _prepare_environment(self, force_reload: bool = False):
85
+ """
86
+ Common setup. Skips if already initialized, unless force_reload is True.
87
+ """
88
+ if self._is_env_ready and not force_reload:
89
+ # Environment is already prepared, skipping setup.
90
+ return
91
+
92
+ # 1. Run Setup
93
+ self.init_general(self.config.get_general_cfg())
94
+ self.prepare_dataset(self.config.get_dataset_cfg())
95
+
96
+ # 2. Update metric backend (refresh if needed)
97
+ self.metric_backend = self.prepare_metrics(self.config.get_metric_cfg())
98
+
99
+ # 3. Mark as ready
100
+ self._is_env_ready = True
101
+
102
+ # -----------------------
103
+ # Main Experiment Runner
104
+ # -----------------------
105
+ def run_exp(self, should_calc_metrics=True, reload_env=False, *args, **kwargs):
106
+ """
107
+ Run the whole experiment pipeline.
108
+ :param reload_env: If True, forces dataset/general init to run again.
109
+ :param should_calc_metrics: Whether to calculate and save metrics after execution.
110
+ :kwargs Params:
111
+ + 'outfile' to save csv file results,
112
+ + 'outdir' to set output directory for experiment results.
113
+ + 'return_df' to return a DataFrame of results instead of a dictionary.
114
+
115
+ Full pipeline:
116
+ 1. Init
117
+ 2. Prepare Environment (General + Dataset + Metrics)
118
+ 3. Save Config
119
+ 4. Execute
120
+ 5. Calculate & Save Metrics
121
+ """
122
+ self._prepare_environment(force_reload=reload_env)
123
+
124
+ # Any pre-exec setup (loading models, etc)
125
+ self.before_exec_exp_once(*args, **kwargs)
126
+ # Save config before running
127
+ self.config.save_to_outdir()
128
+
129
+ # Execute experiment
130
+ results = self.exec_exp(*args, **kwargs)
131
+
132
+ if should_calc_metrics and results is not None:
133
+ metrics_data, extra_data = self._validate_and_unpack(results)
134
+ # Calculate & Save metrics
135
+ perf_results = self.calc_perfs(
136
+ raw_metrics_data=metrics_data, extra_data=extra_data, *args, **kwargs
137
+ )
138
+ return perf_results
139
+ else:
140
+ return results
141
+
142
+ # -----------------------
143
+ # Main Experiment Evaluator
144
+ # -----------------------
145
+ def eval_exp(self, reload_env=False, *args, **kwargs):
146
+ """
147
+ Run evaluation only.
148
+ :param reload_env: If True, forces dataset/general init to run again.
149
+ """
150
+ self._prepare_environment(force_reload=reload_env)
151
+ results = self.exec_eval(*args, **kwargs)
152
+ if results is not None:
153
+ metrics_data, extra_data = self._validate_and_unpack(results)
154
+ return self.calc_perfs(
155
+ raw_metrics_data=metrics_data, extra_data=extra_data, *args, **kwargs
156
+ )
157
+ return None
@@ -0,0 +1,108 @@
1
+ import os
2
+ import yaml
3
+ import numpy as np
4
+ from typing import Dict, Any, List
5
+
6
+ from ...common.common import *
7
+ from ...filetype import yamlfile
8
+
9
+ class ParamGen:
10
+ @staticmethod
11
+ def build_from_file(params_file):
12
+ builder = ParamGen(params_file)
13
+ return builder.params
14
+
15
+ def __init__(self, params_file=None):
16
+ self.params = {}
17
+ assert os.path.isfile(params_file), f"params_file not found: {params_file}"
18
+ self.params = self._build(params_file)
19
+
20
+ def _expand_param(self, param_name: str, config: Dict[str, Any]) -> List[Any]:
21
+ """
22
+ Validates and expands the values for a single parameter configuration.
23
+
24
+ Args:
25
+ param_name: The name of the parameter being processed.
26
+ config: The configuration dictionary for this parameter.
27
+
28
+ Returns:
29
+ A list of the expanded values for the parameter.
30
+
31
+ Raises:
32
+ TypeError: If the configuration or its values have an incorrect type.
33
+ ValueError: If the configuration is missing keys or has an invalid structure.
34
+ """
35
+ # 1. Validate the configuration structure
36
+ if not isinstance(config, dict):
37
+ raise TypeError(f"Config for '{param_name}' must be a dictionary.")
38
+
39
+ if "type" not in config or "values" not in config:
40
+ raise ValueError(
41
+ f"Config for '{param_name}' must contain 'type' and 'values' keys."
42
+ )
43
+
44
+ gen_type = config["type"]
45
+ values = config["values"]
46
+
47
+ # 2. Handle the generation based on type
48
+ if gen_type == "list":
49
+ # Ensure values are returned as a list, even if a single item was provided
50
+ return values if isinstance(values, list) else [values]
51
+
52
+ elif gen_type == "range":
53
+ if not isinstance(values, list) or len(values) != 3:
54
+ raise ValueError(
55
+ f"For 'range' type on '{param_name}', 'values' must be a list of 3 numbers "
56
+ f"[start, end, step], but got: {values}"
57
+ )
58
+
59
+ start, end, step = values
60
+ if all(isinstance(v, int) for v in values):
61
+ return list(range(start, end, step))
62
+ elif all(isinstance(v, (int, float)) for v in values):
63
+ # Use numpy for floating point ranges
64
+ temp_list = list(np.arange(start, end, step))
65
+ # convert to float (not np.float)
66
+ return [float(v) for v in temp_list]
67
+ else:
68
+ raise TypeError(
69
+ f"All 'values' for 'range' on '{param_name}' must be numbers."
70
+ )
71
+
72
+ else:
73
+ raise ValueError(
74
+ f"Invalid 'type' for '{param_name}': '{gen_type}'. Must be 'list' or 'range'."
75
+ )
76
+
77
+ def _build(self, params_file):
78
+ """
79
+ Builds a full optimization configuration by expanding parameter values based on their type.
80
+
81
+ This function processes a dictionary where each key is a parameter name and each value
82
+ is a config dict specifying the 'type' ('list' or 'range') and 'values' for generation.
83
+
84
+ Args:
85
+ opt_cfg: The input configuration dictionary.
86
+ Example:
87
+ {
88
+ "learning_rate": {"type": "range", "values": [0.01, 0.1, 0.01]},
89
+ "optimizer": {"type": "list", "values": ["adam", "sgd"]},
90
+ "epochs": {"type": "list", "values": 100}
91
+ }
92
+
93
+ Returns:
94
+ A dictionary with parameter names mapped to their fully expanded list of values.
95
+ """
96
+ cfg_raw_dict = yamlfile.load_yaml(params_file, to_dict=True)
97
+ if not isinstance(cfg_raw_dict, dict):
98
+ raise TypeError("The entire opt_cfg must be a dictionary.")
99
+
100
+ # Use a dictionary comprehension for a clean and efficient build
101
+ return {
102
+ param_name: self._expand_param(param_name, config)
103
+ for param_name, config in cfg_raw_dict.items()
104
+ }
105
+
106
+ def save(self, outfile):
107
+ with open(outfile, "w") as f:
108
+ yaml.dump(self.params, f)
@@ -0,0 +1,117 @@
1
+ import os
2
+ import glob
3
+ import wandb
4
+ import argparse
5
+ import subprocess
6
+
7
+ from tqdm import tqdm
8
+ from rich.console import Console
9
+
10
+ console = Console()
11
+
12
+ def sync_runs(outdir):
13
+ outdir = os.path.abspath(outdir)
14
+ assert os.path.exists(outdir), f"Output directory {outdir} does not exist."
15
+ sub_dirs = [name for name in os.listdir(outdir) if os.path.isdir(os.path.join(outdir, name))]
16
+ assert len(sub_dirs) > 0, f"No subdirectories found in {outdir}."
17
+ console.rule("Parent Directory")
18
+ console.print(f"[yellow]{outdir}[/yellow]")
19
+
20
+ exp_dirs = [os.path.join(outdir, sub_dir) for sub_dir in sub_dirs]
21
+ wandb_dirs = []
22
+ for exp_dir in exp_dirs:
23
+ wandb_dirs.extend(glob.glob(f"{exp_dir}/wandb/*run-*"))
24
+ if len(wandb_dirs) == 0:
25
+ console.print(f"No wandb runs found in {outdir}.")
26
+ return
27
+ else:
28
+ console.print(f"Found [bold]{len(wandb_dirs)}[/bold] wandb runs in {outdir}.")
29
+ for i, wandb_dir in enumerate(wandb_dirs):
30
+ console.rule(f"Syncing wandb run {i + 1}/{len(wandb_dirs)}")
31
+ console.print(f"Syncing: {wandb_dir}")
32
+ process = subprocess.Popen(
33
+ ["wandb", "sync", wandb_dir],
34
+ stdout=subprocess.PIPE,
35
+ stderr=subprocess.STDOUT,
36
+ text=True,
37
+ )
38
+
39
+ for line in process.stdout:
40
+ console.print(line.strip())
41
+ if " ERROR Error while calling W&B API" in line:
42
+ break
43
+ process.stdout.close()
44
+ process.wait()
45
+ if process.returncode != 0:
46
+ console.print(f"[red]Error syncing {wandb_dir}. Return code: {process.returncode}[/red]")
47
+ else:
48
+ console.print(f"Successfully synced {wandb_dir}.")
49
+
50
+ def delete_runs(project, pattern=None):
51
+ console.rule("Delete W&B Runs")
52
+ confirm_msg = f"Are you sure you want to delete all runs in"
53
+ confirm_msg += f" \n\tproject: [red]{project}[/red]"
54
+ if pattern:
55
+ confirm_msg += f"\n\tpattern: [blue]{pattern}[/blue]"
56
+
57
+ console.print(confirm_msg)
58
+ confirmation = input(f"This action cannot be undone. [y/N]: ").strip().lower()
59
+ if confirmation != "y":
60
+ print("Cancelled.")
61
+ return
62
+
63
+ print("Confirmed. Proceeding...")
64
+ api = wandb.Api()
65
+ runs = api.runs(project)
66
+
67
+ deleted = 0
68
+ console.rule("Deleting W&B Runs")
69
+ if len(runs) == 0:
70
+ print("No runs found in the project.")
71
+ return
72
+ for run in tqdm(runs):
73
+ if pattern is None or pattern in run.name:
74
+ run.delete()
75
+ console.print(f"Deleted run: [red]{run.name}[/red]")
76
+ deleted += 1
77
+
78
+ console.print(f"Total runs deleted: {deleted}")
79
+
80
+
81
+ def valid_argument(args):
82
+ if args.op == "sync":
83
+ assert os.path.exists(args.outdir), f"Output directory {args.outdir} does not exist."
84
+ elif args.op == "delete":
85
+ assert isinstance(args.project, str) and len(args.project.strip()) > 0, "Project name must be a non-empty string."
86
+ else:
87
+ raise ValueError(f"Unknown operation: {args.op}")
88
+
89
+ def parse_args():
90
+ parser = argparse.ArgumentParser(description="Operations on W&B runs")
91
+ parser.add_argument("-op", "--op", type=str, help="Operation to perform", default="sync", choices=["delete", "sync"])
92
+ parser.add_argument("-prj", "--project", type=str, default="fire-paper2-2025", help="W&B project name")
93
+ parser.add_argument("-outdir", "--outdir", type=str, help="arg1 description", default="./zout/train")
94
+ parser.add_argument("-pt", "--pattern",
95
+ type=str,
96
+ default=None,
97
+ help="Run name pattern to match for deletion",
98
+ )
99
+
100
+ return parser.parse_args()
101
+
102
+
103
+ def main():
104
+ args = parse_args()
105
+ # Validate arguments, stop if invalid
106
+ valid_argument(args)
107
+
108
+ op = args.op
109
+ if op == "sync":
110
+ sync_runs(args.outdir)
111
+ elif op == "delete":
112
+ delete_runs(args.project, args.pattern)
113
+ else:
114
+ raise ValueError(f"Unknown operation: {op}")
115
+
116
+ if __name__ == "__main__":
117
+ main()
File without changes
@@ -0,0 +1,41 @@
1
+ import yaml
2
+ from typing import Any
3
+
4
+ from rich.pretty import pprint
5
+ from dataclasses import make_dataclass
6
+
7
+ from ...filetype import yamlfile
8
+
9
+ def dict_to_dataclass(name: str, data: dict):
10
+ fields = []
11
+ values = {}
12
+
13
+ for key, value in data.items():
14
+ if isinstance(value, dict):
15
+ sub_dc = dict_to_dataclass(key.capitalize(), value)
16
+ fields.append((key, type(sub_dc)))
17
+ values[key] = sub_dc
18
+ else:
19
+ field_type = type(value) if value is not None else Any
20
+ fields.append((key, field_type))
21
+ values[key] = value
22
+
23
+ DC = make_dataclass(name.capitalize(), fields)
24
+ return DC(**values)
25
+
26
+ def yaml_to_dataclass(name: str, yaml_str: str):
27
+ data = yaml.safe_load(yaml_str)
28
+ return dict_to_dataclass(name, data)
29
+
30
+
31
+ def yamlfile_to_dataclass(name: str, file_path: str):
32
+ data_dict = yamlfile.load_yaml(file_path, to_dict=True)
33
+ if "__base__" in data_dict:
34
+ del data_dict["__base__"]
35
+ return dict_to_dataclass(name, data_dict)
36
+
37
+ if __name__ == "__main__":
38
+ cfg = yamlfile_to_dataclass("Config", "test/dataclass_util_test_cfg.yaml")
39
+
40
+ # ! NOTICE: after print out this dataclass, we can copy the output and paste it into CHATGPT to generate a list of needed dataclass classes using `from dataclass_wizard import YAMLWizard`
41
+ pprint(cfg)
@@ -0,0 +1,208 @@
1
+ # This script create a test version
2
+ # of the watcam (wc) dataset
3
+ # for testing the tflite model
4
+
5
+ from argparse import ArgumentParser
6
+
7
+ import os
8
+ import click
9
+ import shutil
10
+ from tqdm import tqdm
11
+ from rich import inspect
12
+ from rich.pretty import pprint
13
+ from torchvision.datasets import ImageFolder
14
+ from sklearn.model_selection import StratifiedShuffleSplit, ShuffleSplit
15
+
16
+ from ...common.common import console, seed_everything, ConsoleLog
17
+ from ...system import filesys as fs
18
+
19
+ def parse_args():
20
+ parser = ArgumentParser(description="desc text")
21
+ parser.add_argument(
22
+ "-indir",
23
+ "--indir",
24
+ type=str,
25
+ help="orignal dataset path",
26
+ )
27
+ parser.add_argument(
28
+ "-outdir",
29
+ "--outdir",
30
+ type=str,
31
+ help="dataset out path",
32
+ default=".", # default to current dir
33
+ )
34
+ parser.add_argument(
35
+ "-val_size",
36
+ "--val_size",
37
+ type=float,
38
+ help="validation size", # no default value to force user to input
39
+ default=0.2,
40
+ )
41
+ # add using StratifiedShuffleSplit or ShuffleSplit
42
+ parser.add_argument(
43
+ "-seed",
44
+ "--seed",
45
+ type=int,
46
+ help="random seed",
47
+ default=42,
48
+ )
49
+ parser.add_argument(
50
+ "-inplace",
51
+ "--inplace",
52
+ action="store_true",
53
+ help="inplace operation, will overwrite the outdir if exists",
54
+ )
55
+
56
+ parser.add_argument(
57
+ "-stratified",
58
+ "--stratified",
59
+ action="store_true",
60
+ help="use StratifiedShuffleSplit instead of ShuffleSplit",
61
+ )
62
+ parser.add_argument(
63
+ "-no_train",
64
+ "--no_train",
65
+ action="store_true",
66
+ help="only create test set, no train set",
67
+ )
68
+ parser.add_argument(
69
+ "-reverse",
70
+ "--reverse",
71
+ action="store_true",
72
+ help="combine train and val set back to original dataset",
73
+ )
74
+ return parser.parse_args()
75
+
76
+
77
+ def move_images(image_paths, target_set_dir):
78
+ for img_path in tqdm(image_paths):
79
+ # get folder name of the image
80
+ img_dir = os.path.dirname(img_path)
81
+ out_cls_dir = os.path.join(target_set_dir, os.path.basename(img_dir))
82
+ if not os.path.exists(out_cls_dir):
83
+ os.makedirs(out_cls_dir)
84
+ # move the image to the class folder
85
+ shutil.move(img_path, out_cls_dir)
86
+
87
+
88
+ def split_dataset_cls(
89
+ indir, outdir, val_size, seed, inplace, stratified_split, no_train
90
+ ):
91
+ seed_everything(seed)
92
+ console.rule("Config confirm?")
93
+ pprint(locals())
94
+ click.confirm("Continue?", abort=True)
95
+ assert os.path.exists(indir), f"{indir} does not exist"
96
+
97
+ if not inplace:
98
+ assert (not inplace) and (
99
+ not os.path.exists(outdir)
100
+ ), f"{outdir} already exists; SKIP ...."
101
+
102
+ if inplace:
103
+ outdir = indir
104
+ if not os.path.exists(outdir):
105
+ os.makedirs(outdir)
106
+
107
+ console.rule(f"Creating train/val dataset")
108
+
109
+ sss = (
110
+ ShuffleSplit(n_splits=1, test_size=val_size)
111
+ if not stratified_split
112
+ else StratifiedShuffleSplit(n_splits=1, test_size=val_size)
113
+ )
114
+
115
+ pprint({"split strategy": sss, "indir": indir, "outdir": outdir})
116
+ dataset = ImageFolder(
117
+ root=indir,
118
+ transform=None,
119
+ )
120
+ train_dataset_indices = None
121
+ val_dataset_indices = None # val here means test
122
+ for train_indices, val_indices in sss.split(dataset.samples, dataset.targets):
123
+ train_dataset_indices = train_indices
124
+ val_dataset_indices = val_indices
125
+
126
+ # get image paths for train/val split dataset
127
+ train_image_paths = [dataset.imgs[i][0] for i in train_dataset_indices]
128
+ val_image_paths = [dataset.imgs[i][0] for i in val_dataset_indices]
129
+
130
+ # start creating train/val folders then move images
131
+ out_train_dir = os.path.join(outdir, "train")
132
+ out_val_dir = os.path.join(outdir, "val")
133
+ if inplace:
134
+ assert os.path.exists(out_train_dir) == False, f"{out_train_dir} already exists"
135
+ assert os.path.exists(out_val_dir) == False, f"{out_val_dir} already exists"
136
+
137
+ os.makedirs(out_train_dir)
138
+ os.makedirs(out_val_dir)
139
+
140
+ if not no_train:
141
+ with ConsoleLog(f"Moving train images to {out_train_dir} "):
142
+ move_images(train_image_paths, out_train_dir)
143
+ else:
144
+ pprint("test only, skip moving train images")
145
+ # remove out_train_dir
146
+ shutil.rmtree(out_train_dir)
147
+
148
+ with ConsoleLog(f"Moving val images to {out_val_dir} "):
149
+ move_images(val_image_paths, out_val_dir)
150
+
151
+ if inplace:
152
+ pprint(f"remove all folders, except train and val")
153
+ for cls_dir in os.listdir(outdir):
154
+ if cls_dir not in ["train", "val"]:
155
+ shutil.rmtree(os.path.join(indir, cls_dir))
156
+
157
+
158
+ def reverse_split_ds(indir):
159
+ console.rule(f"Reversing split dataset <{indir}>...")
160
+ ls_dirs = os.listdir(indir)
161
+ # make sure there are only two dirs 'train' and 'val'
162
+ assert len(ls_dirs) == 2, f"Found more than 2 dirs: {len(ls_dirs) } dirs"
163
+ assert "train" in ls_dirs, f"train dir not found in {indir}"
164
+ assert "val" in ls_dirs, f"val dir not found in {indir}"
165
+ train_dir = os.path.join(indir, "train")
166
+ val_dir = os.path.join(indir, "val")
167
+ all_train_files = fs.filter_files_by_extension(
168
+ train_dir, ["jpg", "jpeg", "png", "bmp", "gif", "tiff"]
169
+ )
170
+ all_val_files = fs.filter_files_by_extension(
171
+ val_dir, ["jpg", "jpeg", "png", "bmp", "gif", "tiff"]
172
+ )
173
+ # move all files from train to indir
174
+ with ConsoleLog(f"Moving train images to {indir} "):
175
+ move_images(all_train_files, indir)
176
+ with ConsoleLog(f"Moving val images to {indir} "):
177
+ move_images(all_val_files, indir)
178
+ with ConsoleLog(f"Removing train and val dirs"):
179
+ # remove train and val dirs
180
+ shutil.rmtree(train_dir)
181
+ shutil.rmtree(val_dir)
182
+
183
+
184
+ def main():
185
+ args = parse_args()
186
+ indir = args.indir
187
+ outdir = args.outdir
188
+ if outdir == ".":
189
+ # get current folder of the indir
190
+ indir_parent_dir = os.path.dirname(os.path.normpath(indir))
191
+ indir_name = os.path.basename(indir)
192
+ outdir = os.path.join(indir_parent_dir, f"{indir_name}_split")
193
+ val_size = args.val_size
194
+ seed = args.seed
195
+ inplace = args.inplace
196
+ stratified_split = args.stratified
197
+ no_train = args.no_train
198
+ reverse = args.reverse
199
+ if not reverse:
200
+ split_dataset_cls(
201
+ indir, outdir, val_size, seed, inplace, stratified_split, no_train
202
+ )
203
+ else:
204
+ reverse_split_ds(indir)
205
+
206
+
207
+ if __name__ == "__main__":
208
+ main()