halib 0.1.7__py3-none-any.whl → 0.1.99__py3-none-any.whl

halib/research/dataset.py

@@ -0,0 +1,208 @@
+# This script create a test version
+# of the watcam (wc) dataset
+# for testing the tflite model
+
+from argparse import ArgumentParser
+
+import os
+import click
+import shutil
+from tqdm import tqdm
+from rich import inspect
+from rich.pretty import pprint
+from torchvision.datasets import ImageFolder
+from sklearn.model_selection import StratifiedShuffleSplit, ShuffleSplit
+
+from ..system import filesys as fs
+from ..common import console, seed_everything, ConsoleLog
+
+def parse_args():
+    parser = ArgumentParser(description="desc text")
+    parser.add_argument(
+        "-indir",
+        "--indir",
+        type=str,
+        help="orignal dataset path",
+    )
+    parser.add_argument(
+        "-outdir",
+        "--outdir",
+        type=str,
+        help="dataset out path",
+        default=".",  # default to current dir
+    )
+    parser.add_argument(
+        "-val_size",
+        "--val_size",
+        type=float,
+        help="validation size",  # no default value to force user to input
+        default=0.2,
+    )
+    # add using StratifiedShuffleSplit or ShuffleSplit
+    parser.add_argument(
+        "-seed",
+        "--seed",
+        type=int,
+        help="random seed",
+        default=42,
+    )
+    parser.add_argument(
+        "-inplace",
+        "--inplace",
+        action="store_true",
+        help="inplace operation, will overwrite the outdir if exists",
+    )
+
+    parser.add_argument(
+        "-stratified",
+        "--stratified",
+        action="store_true",
+        help="use StratifiedShuffleSplit instead of ShuffleSplit",
+    )
+    parser.add_argument(
+        "-no_train",
+        "--no_train",
+        action="store_true",
+        help="only create test set, no train set",
+    )
+    parser.add_argument(
+        "-reverse",
+        "--reverse",
+        action="store_true",
+        help="combine train and val set back to original dataset",
+    )
+    return parser.parse_args()
+
+
+def move_images(image_paths, target_set_dir):
+    for img_path in tqdm(image_paths):
+        # get folder name of the image
+        img_dir = os.path.dirname(img_path)
+        out_cls_dir = os.path.join(target_set_dir, os.path.basename(img_dir))
+        if not os.path.exists(out_cls_dir):
+            os.makedirs(out_cls_dir)
+        # move the image to the class folder
+        shutil.move(img_path, out_cls_dir)
+
+
+def split_dataset_cls(
+    indir, outdir, val_size, seed, inplace, stratified_split, no_train
+):
+    seed_everything(seed)
+    console.rule("Config confirm?")
+    pprint(locals())
+    click.confirm("Continue?", abort=True)
+    assert os.path.exists(indir), f"{indir} does not exist"
+
+    if not inplace:
+        assert (not inplace) and (
+            not os.path.exists(outdir)
+        ), f"{outdir} already exists; SKIP ...."
+
+    if inplace:
+        outdir = indir
+    if not os.path.exists(outdir):
+        os.makedirs(outdir)
+
+    console.rule(f"Creating train/val dataset")
+
+    sss = (
+        ShuffleSplit(n_splits=1, test_size=val_size)
+        if not stratified_split
+        else StratifiedShuffleSplit(n_splits=1, test_size=val_size)
+    )
+
+    pprint({"split strategy": sss, "indir": indir, "outdir": outdir})
+    dataset = ImageFolder(
+        root=indir,
+        transform=None,
+    )
+    train_dataset_indices = None
+    val_dataset_indices = None  # val here means test
+    for train_indices, val_indices in sss.split(dataset.samples, dataset.targets):
+        train_dataset_indices = train_indices
+        val_dataset_indices = val_indices
+
+    # get image paths for train/val split dataset
+    train_image_paths = [dataset.imgs[i][0] for i in train_dataset_indices]
+    val_image_paths = [dataset.imgs[i][0] for i in val_dataset_indices]
+
+    # start creating train/val folders then move images
+    out_train_dir = os.path.join(outdir, "train")
+    out_val_dir = os.path.join(outdir, "val")
+    if inplace:
+        assert os.path.exists(out_train_dir) == False, f"{out_train_dir} already exists"
+        assert os.path.exists(out_val_dir) == False, f"{out_val_dir} already exists"
+
+    os.makedirs(out_train_dir)
+    os.makedirs(out_val_dir)
+
+    if not no_train:
+        with ConsoleLog(f"Moving train images to {out_train_dir} "):
+            move_images(train_image_paths, out_train_dir)
+    else:
+        pprint("test only, skip moving train images")
+        # remove out_train_dir
+        shutil.rmtree(out_train_dir)
+
+    with ConsoleLog(f"Moving val images to {out_val_dir} "):
+        move_images(val_image_paths, out_val_dir)
+
+    if inplace:
+        pprint(f"remove all folders, except train and val")
+        for cls_dir in os.listdir(outdir):
+            if cls_dir not in ["train", "val"]:
+                shutil.rmtree(os.path.join(indir, cls_dir))
+
+
+def reverse_split_ds(indir):
+    console.rule(f"Reversing split dataset <{indir}>...")
+    ls_dirs = os.listdir(indir)
+    # make sure there are only two dirs 'train' and 'val'
+    assert len(ls_dirs) == 2, f"Found more than 2 dirs: {len(ls_dirs) } dirs"
+    assert "train" in ls_dirs, f"train dir not found in {indir}"
+    assert "val" in ls_dirs, f"val dir not found in {indir}"
+    train_dir = os.path.join(indir, "train")
+    val_dir = os.path.join(indir, "val")
+    all_train_files = fs.filter_files_by_extension(
+        train_dir, ["jpg", "jpeg", "png", "bmp", "gif", "tiff"]
+    )
+    all_val_files = fs.filter_files_by_extension(
+        val_dir, ["jpg", "jpeg", "png", "bmp", "gif", "tiff"]
+    )
+    # move all files from train to indir
+    with ConsoleLog(f"Moving train images to {indir} "):
+        move_images(all_train_files, indir)
+    with ConsoleLog(f"Moving val images to {indir} "):
+        move_images(all_val_files, indir)
+    with ConsoleLog(f"Removing train and val dirs"):
+        # remove train and val dirs
+        shutil.rmtree(train_dir)
+        shutil.rmtree(val_dir)
+
+
+def main():
+    args = parse_args()
+    indir = args.indir
+    outdir = args.outdir
+    if outdir == ".":
+        # get current folder of the indir
+        indir_parent_dir = os.path.dirname(os.path.normpath(indir))
+        indir_name = os.path.basename(indir)
+        outdir = os.path.join(indir_parent_dir, f"{indir_name}_split")
+    val_size = args.val_size
+    seed = args.seed
+    inplace = args.inplace
+    stratified_split = args.stratified
+    no_train = args.no_train
+    reverse = args.reverse
+    if not reverse:
+        split_dataset_cls(
+            indir, outdir, val_size, seed, inplace, stratified_split, no_train
+        )
+    else:
+        reverse_split_ds(indir)
+
+
+if __name__ == "__main__":
+    main()

halib/research/flop_csv.py

@@ -0,0 +1,34 @@
+from halib import *
+from flops import _calculate_flops_for_model
+
+from halib import *
+from argparse import ArgumentParser
+
+
+def main():
+    csv_file = "./results-imagenet.csv"
+    df = pd.read_csv(csv_file)
+    # make param_count column as float
+    # df['param_count'] = df['param_count'].astype(float)
+    df['param_count'] = pd.to_numeric(df['param_count'], errors='coerce').fillna(99999).astype(float)
+    df = df[df['param_count'] < 5.0]  # filter models with param_count < 20M
+
+    dict_ls = []
+
+    for index, row in tqdm(df.iterrows()):
+        console.rule(f"Row {index+1}/{len(df)}")
+        model = row['model']
+        num_class = 2
+        _, _, mflops = _calculate_flops_for_model(model, num_class)
+        dict_ls.append({'model': model, 'param_count': row['param_count'], 'mflops': mflops})
+
+    # Create a DataFrame from the list of dictionaries
+    result_df = pd.DataFrame(dict_ls)
+
+    final_df = pd.merge(df, result_df, on=['model', 'param_count'])
+    final_df.sort_values(by='mflops', inplace=True, ascending=True)
+    csvfile.fn_display_df(final_df)
+
+
+if __name__ == "__main__":
+    main()

halib/research/flops.py

@@ -0,0 +1,156 @@
+import os
+import sys
+import torch
+import timm
+from argparse import ArgumentParser
+from fvcore.nn import FlopCountAnalysis
+from halib import *
+from halib.filetype import csvfile
+from curriculum.utils.config import *
+from curriculum.utils.model_helper import *
+
+
+# ---------------------------------------------------------------------
+# Argument Parser
+# ---------------------------------------------------------------------
+def parse_args():
+    parser = ArgumentParser(description="Calculate FLOPs for TIMM or trained models")
+
+    # Option 1: Direct TIMM model
+    parser.add_argument(
+        "--model_name", type=str, help="TIMM model name (e.g., efficientnet_b0)"
+    )
+    parser.add_argument(
+        "--num_classes", type=int, default=1000, help="Number of output classes"
+    )
+
+    # Option 2: Experiment directory
+    parser.add_argument(
+        "--indir",
+        type=str,
+        default=None,
+        help="Directory containing trained experiment (with .yaml and .pth)",
+    )
+    parser.add_argument(
+        "-o", "--o", action="store_true", help="Open output CSV after saving"
+    )
+    return parser.parse_args()
+
+
+# ---------------------------------------------------------------------
+# Helper Functions
+# ---------------------------------------------------------------------
+def _get_list_of_proc_dirs(indir):
+    assert os.path.exists(indir), f"Input directory {indir} does not exist."
+    pth_files = [f for f in os.listdir(indir) if f.endswith(".pth")]
+    if len(pth_files) > 0:
+        return [indir]
+    return [
+        os.path.join(indir, f)
+        for f in os.listdir(indir)
+        if os.path.isdir(os.path.join(indir, f))
+    ]
+
+
+def _calculate_flops_for_model(model_name, num_classes):
+    """Calculate FLOPs for a plain TIMM model."""
+    try:
+        model = timm.create_model(model_name, pretrained=False, num_classes=num_classes)
+        input_size = timm.data.resolve_data_config(model.default_cfg)["input_size"]
+        dummy_input = torch.randn(1, *input_size)
+        model.eval() # ! set to eval mode to avoid some warnings or errors
+        flops = FlopCountAnalysis(model, dummy_input)
+        gflops = flops.total() / 1e9
+        mflops = flops.total() / 1e6
+        print(f"\nModel: **{model_name}**, Classes: {num_classes}")
+        print(f"Input size: {input_size}, FLOPs: **{gflops:.3f} GFLOPs**, **{mflops:.3f} MFLOPs**\n")
+        return model_name, gflops, mflops
+    except Exception as e:
+        print(f"[Error] Could not calculate FLOPs for {model_name}: {e}")
+        return model_name, -1, -1
+
+
+def _calculate_flops_for_experiment(exp_dir):
+    """Calculate FLOPs for a trained experiment directory."""
+    yaml_files = [f for f in os.listdir(exp_dir) if f.endswith(".yaml")]
+    pth_files = [f for f in os.listdir(exp_dir) if f.endswith(".pth")]
+
+    assert (
+        len(yaml_files) == 1
+    ), f"Expected 1 YAML file in {exp_dir}, found {len(yaml_files)}"
+    assert (
+        len(pth_files) == 1
+    ), f"Expected 1 PTH file in {exp_dir}, found {len(pth_files)}"
+
+    exp_cfg_yaml = os.path.join(exp_dir, yaml_files[0])
+    cfg = ExpConfig.from_yaml(exp_cfg_yaml)
+    ds_label_list = cfg.dataset.get_label_list()
+
+    try:
+        model = build_model(
+            cfg.model.name, num_classes=len(ds_label_list), pretrained=True
+        )
+        model_weights_path = os.path.join(exp_dir, pth_files[0])
+        model.load_state_dict(torch.load(model_weights_path, map_location="cpu"))
+        model.eval()
+
+        input_size = timm.data.resolve_data_config(model.default_cfg)["input_size"]
+        dummy_input = torch.randn(1, *input_size)
+        flops = FlopCountAnalysis(model, dummy_input)
+        gflops = flops.total() / 1e9
+        mflops = flops.total() / 1e6
+
+        return str(cfg), cfg.model.name, gflops, mflops
+    except Exception as e:
+        console.print(f"[red] Error processing {exp_dir}: {e}[/red]")
+        return str(cfg), cfg.model.name, -1, -1
+
+
+# ---------------------------------------------------------------------
+# Main Entry
+# ---------------------------------------------------------------------
+def main():
+    args = parse_args()
+
+    # Case 1: Direct TIMM model input
+    if args.model_name:
+        _calculate_flops_for_model(args.model_name, args.num_classes)
+        return
+
+    # Case 2: Experiment directory input
+    if args.indir is None:
+        print("[Error] Either --model_name or --indir must be specified.")
+        return
+
+    proc_dirs = _get_list_of_proc_dirs(args.indir)
+    pprint(proc_dirs)
+
+    dfmk = csvfile.DFCreator()
+    TABLE_NAME = "model_flops_results"
+    dfmk.create_table(TABLE_NAME, ["exp_name", "model_name", "gflops", "mflops"])
+
+    console.rule(f"Calculating FLOPs for models in {len(proc_dirs)} dir(s)...")
+    rows = []
+    for exp_dir in tqdm(proc_dirs):
+        dir_name = os.path.basename(exp_dir)
+        console.rule(f"{dir_name}")
+        exp_name, model_name, gflops, mflops = _calculate_flops_for_experiment(exp_dir)
+        rows.append([exp_name, model_name, gflops, mflops])
+
+    dfmk.insert_rows(TABLE_NAME, rows)
+    dfmk.fill_table_from_row_pool(TABLE_NAME)
+
+    outfile = f"zout/zreport/{now_str()}_model_flops_results.csv"
+    dfmk[TABLE_NAME].to_csv(outfile, sep=";", index=False)
+    csvfile.fn_display_df(dfmk[TABLE_NAME])
+
+    if args.o:
+        os.system(f"start {outfile}")
+
+
+# ---------------------------------------------------------------------
+# Script Entry
+# ---------------------------------------------------------------------
+if __name__ == "__main__":
+    sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+    main()

halib/research/metrics.py

@@ -0,0 +1,133 @@
+# -------------------------------
+# Metrics Backend Interface
+# -------------------------------
+import inspect
+from typing import Dict, Union, List, Any
+from abc import ABC, abstractmethod
+
+class MetricsBackend(ABC):
+    """Interface for pluggable metrics computation backends."""
+
+    def __init__(self, metrics_info: Union[List[str], Dict[str, Any]]):
+        """
+        Initialize the backend with optional metrics_info.
+        """
+        self.metric_info = metrics_info
+        self.validate_metrics_info(self.metric_info)
+
+    @property
+    def metric_names(self) -> List[str]:
+        """
+        Return a list of metric names.
+        If metric_info is a dict, return its keys; if it's a list, return it directly.
+        """
+        if isinstance(self.metric_info, dict):
+            return list(self.metric_info.keys())
+        elif isinstance(self.metric_info, list):
+            return self.metric_info
+        else:
+            raise TypeError("metric_info must be a list or a dict")
+
+    def validate_metrics_info(self, metrics_info):
+        if isinstance(metrics_info, list):
+            return metrics_info
+        elif isinstance(metrics_info, dict):
+            return {k: v for k, v in metrics_info.items() if isinstance(k, str)}
+        else:
+            raise TypeError(
+                "metrics_info must be a list of strings or a dict with string keys"
+            )
+
+    @abstractmethod
+    def compute_metrics(
+        self, metrics_info: Union[List[str], Dict[str, Any]],  metrics_data_dict: Dict[str, Any], *args, **kwargs
+    ) -> Dict[str, Any]:
+        pass
+
+    def prepare_metrics_backend_data(
+        self, raw_metric_data, *args, **kwargs
+    ):
+        """
+        Prepare the data for the metrics backend.
+        This method can be overridden by subclasses to customize data preparation.
+        """
+        return raw_metric_data
+
+    def calc_metrics(
+        self,  metrics_data_dict: Dict[str, Any], *args, **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Calculate metrics based on the provided metrics_info and data.
+        This method should be overridden by subclasses to implement specific metric calculations.
+        """
+        # prevalidate the metrics_data_dict
+        for metric in self.metric_names:
+            if metric not in metrics_data_dict:
+                raise ValueError(f"Metric '{metric}' not found in provided data.")
+        # Prepare the data for the backend
+        metrics_data_dict = self.prepare_metrics_backend_data(
+            metrics_data_dict, *args, **kwargs
+        )
+        # Call the abstract method to compute metrics
+        return self.compute_metrics(self.metric_info, metrics_data_dict, *args, **kwargs)
+
+class TorchMetricsBackend(MetricsBackend):
+    """TorchMetrics-based backend implementation."""
+
+    def __init__(self, metrics_info: Union[List[str], Dict[str, Any]]):
+        try:
+            import torch
+            from torchmetrics import Metric
+        except ImportError:
+            raise ImportError(
+                "TorchMetricsBackend requires torch and torchmetrics to be installed."
+            )
+        self.metric_info = metrics_info
+        self.torch = torch
+        self.Metric = Metric
+        self.validate_metrics_info(metrics_info)
+
+    def validate_metrics_info(self, metrics_info):
+        if not isinstance(metrics_info, dict):
+            raise TypeError(
+                "TorchMetricsBackend requires metrics_info as a dict {name: MetricInstance}"
+            )
+        for k, v in metrics_info.items():
+            if not isinstance(k, str):
+                raise TypeError(f"Key '{k}' is not a string")
+            if not isinstance(v, self.Metric):
+                raise TypeError(f"Value for key '{k}' must be a torchmetrics.Metric")
+        return metrics_info
+
+    def compute_metrics(self, metrics_info, metrics_data_dict, *args, **kwargs):
+        out_dict = {}
+        for metric, metric_instance in metrics_info.items():
+            if metric not in metrics_data_dict:
+                raise ValueError(f"Metric '{metric}' not found in provided data.")
+
+            metric_data = metrics_data_dict[metric]
+            sig = inspect.signature(metric_instance.update)
+            expected_args = list(sig.parameters.values())
+
+            if isinstance(metric_data, dict):
+                args = [metric_data[param.name] for param in expected_args]
+            elif isinstance(metric_data, (list, tuple)):
+                args = metric_data
+            else:
+                args = metric_data
+            if len(expected_args) == 1:
+                metric_instance.update(args)
+            else:
+                metric_instance.update(*args)
+
+            computed_value = metric_instance.compute()
+            if isinstance(computed_value, self.torch.Tensor):
+                computed_value = (
+                    computed_value.item()
+                    if computed_value.numel() == 1
+                    else computed_value.tolist()
+                )
+
+
+            out_dict[metric] = computed_value
+        return out_dict

halib/research/mics.py

@@ -0,0 +1,68 @@
+from ..common import *
+from ..filetype import csvfile
+import pandas as pd
+import platform
+
+
+PC_NAME_TO_ABBR = {
+    "DESKTOP-JQD9K01": "MainPC",
+    "DESKTOP-5IRHU87": "MSI_Laptop",
+    "DESKTOP-96HQCNO": "4090_SV",
+    "DESKTOP-Q2IKLC0": "4GPU_SV",
+    "DESKTOP-QNS3DNF": "1GPU_SV"
+}
+
+DEFAULT_ABBR_WORKING_DISK = {
+    "MainPC": "E:",
+    "MSI_Laptop": "D:",
+    "4090_SV": "E:",
+    "4GPU_SV": "D:",
+}
+
+def list_PCs(show=True):
+    df = pd.DataFrame(list(PC_NAME_TO_ABBR.items()), columns=["PC Name", "Abbreviation"])
+    if show:
+        csvfile.fn_display_df(df)
+    return df
+
+def get_PC_name():
+    return platform.node()
+
+def get_PC_abbr_name():
+    pc_name = get_PC_name()
+    return PC_NAME_TO_ABBR.get(pc_name, "Unknown")
+
+# ! This funcction search for full paths in the obj and normalize them according to the current platform and working disk
+# ! E.g: "E:/zdataset/DFire", but working_disk: "D:", current_platform: "windows" => "D:/zdataset/DFire"
+# ! E.g: "E:/zdataset/DFire", but working_disk: "D:", current_platform: "linux" => "/mnt/d/zdataset/DFire"
+def normalize_paths(obj, working_disk, current_platform):
+    if isinstance(obj, dict):
+        for key, value in obj.items():
+            obj[key] = normalize_paths(value, working_disk, current_platform)
+        return obj
+    elif isinstance(obj, list):
+        for i, item in enumerate(obj):
+            obj[i] = normalize_paths(item, working_disk, current_platform)
+        return obj
+    elif isinstance(obj, str):
+        # Normalize backslashes to forward slashes for consistency
+        obj = obj.replace("\\", "/")
+        # Regex for Windows-style path: e.g., "E:/zdataset/DFire"
+        win_match = re.match(r"^([A-Z]):/(.*)$", obj)
+        # Regex for Linux-style path: e.g., "/mnt/e/zdataset/DFire"
+        lin_match = re.match(r"^/mnt/([a-z])/(.*)$", obj)
+        if win_match or lin_match:
+            rest = win_match.group(2) if win_match else lin_match.group(2)
+            if current_platform == "windows":
+                # working_disk is like "D:", so "D:/" + rest
+                new_path = working_disk + "/" + rest
+            elif current_platform == "linux":
+                # Extract drive letter from working_disk (e.g., "D:" -> "d")
+                drive_letter = working_disk[0].lower()
+                new_path = "/mnt/" + drive_letter + "/" + rest
+            else:
+                # Unknown platform, return original
+                return obj
+            return new_path
+    # For non-strings or non-path strings, return as is
+    return obj