segment-everything 0.1.0__py3-none-any.whl

segment_everything/vendored/efficientvit/apps/utils/lr.py

@@ -0,0 +1,52 @@
+# EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
+# Han Cai, Junyan Li, Muyan Hu, Chuang Gan, Song Han
+# International Conference on Computer Vision (ICCV), 2023
+
+import math
+
+import torch
+
+from ...models.utils.list import val2list
+
+__all__ = ["CosineLRwithWarmup"]
+
+
+class CosineLRwithWarmup(torch.optim.lr_scheduler._LRScheduler):
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        warmup_steps: int,
+        warmup_lr: float,
+        decay_steps: int or list[int],
+        last_epoch: int = -1,
+    ) -> None:
+        self.warmup_steps = warmup_steps
+        self.warmup_lr = warmup_lr
+        self.decay_steps = val2list(decay_steps)
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self) -> list[float]:
+        if self.last_epoch < self.warmup_steps:
+            return [
+                (base_lr - self.warmup_lr)
+                * (self.last_epoch + 1)
+                / self.warmup_steps
+                + self.warmup_lr
+                for base_lr in self.base_lrs
+            ]
+        else:
+            current_steps = self.last_epoch - self.warmup_steps
+            decay_steps = [0] + self.decay_steps
+            idx = len(decay_steps) - 2
+            for i, decay_step in enumerate(decay_steps[:-1]):
+                if decay_step <= current_steps < decay_steps[i + 1]:
+                    idx = i
+                    break
+            current_steps -= decay_steps[idx]
+            decay_step = decay_steps[idx + 1] - decay_steps[idx]
+            return [
+                0.5
+                * base_lr
+                * (1 + math.cos(math.pi * current_steps / decay_step))
+                for base_lr in self.base_lrs
+            ]

segment_everything/vendored/efficientvit/apps/utils/metric.py

@@ -0,0 +1,43 @@
+# EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
+# Han Cai, Junyan Li, Muyan Hu, Chuang Gan, Song Han
+# International Conference on Computer Vision (ICCV), 2023
+
+import torch
+
+from ..utils.dist import sync_tensor
+
+__all__ = ["AverageMeter"]
+
+
+class AverageMeter:
+    """Computes and stores the average and current value."""
+
+    def __init__(self, is_distributed=True):
+        self.is_distributed = is_distributed
+        self.sum = 0
+        self.count = 0
+
+    def _sync(
+        self, val: torch.Tensor or int or float
+    ) -> torch.Tensor or int or float:
+        return sync_tensor(val, reduce="sum") if self.is_distributed else val
+
+    def update(self, val: torch.Tensor or int or float, delta_n=1):
+        self.count += self._sync(delta_n)
+        self.sum += self._sync(val * delta_n)
+
+    def get_count(self) -> torch.Tensor or int or float:
+        return (
+            self.count.item()
+            if isinstance(self.count, torch.Tensor) and self.count.numel() == 1
+            else self.count
+        )
+
+    @property
+    def avg(self):
+        avg = -1 if self.count == 0 else self.sum / self.count
+        return (
+            avg.item()
+            if isinstance(avg, torch.Tensor) and avg.numel() == 1
+            else avg
+        )

segment_everything/vendored/efficientvit/apps/utils/misc.py

@@ -0,0 +1,101 @@
+# EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
+# Han Cai, Junyan Li, Muyan Hu, Chuang Gan, Song Han
+# International Conference on Computer Vision (ICCV), 2023
+
+import os
+
+import yaml
+
+__all__ = [
+    "parse_with_yaml",
+    "parse_unknown_args",
+    "partial_update_config",
+    "resolve_and_load_config",
+    "load_config",
+    "dump_config",
+]
+
+
+def parse_with_yaml(config_str: str) -> str or dict:
+    try:
+        # add space manually for dict
+        if "{" in config_str and "}" in config_str and ":" in config_str:
+            out_str = config_str.replace(":", ": ")
+        else:
+            out_str = config_str
+        return yaml.safe_load(out_str)
+    except ValueError:
+        # return raw string if parsing fails
+        return config_str
+
+
+def parse_unknown_args(unknown: list) -> dict:
+    """Parse unknown args."""
+    index = 0
+    parsed_dict = {}
+    while index < len(unknown):
+        key, val = unknown[index], unknown[index + 1]
+        index += 2
+        if not key.startswith("--"):
+            continue
+        key = key[2:]
+
+        # try parsing with either dot notation or full yaml notation
+        # Note that the vanilla case "--key value" will be parsed the same
+        if "." in key:
+            # key == a.b.c, val == val --> parsed_dict[a][b][c] = val
+            keys = key.split(".")
+            dict_to_update = parsed_dict
+            for key in keys[:-1]:
+                if not (key in dict_to_update and isinstance(dict_to_update[key], dict)):
+                    dict_to_update[key] = {}
+                dict_to_update = dict_to_update[key]
+            dict_to_update[keys[-1]] = parse_with_yaml(val)  # so we can parse lists, bools, etc...
+        else:
+            parsed_dict[key] = parse_with_yaml(val)
+    return parsed_dict
+
+
+def partial_update_config(config: dict, partial_config: dict) -> dict:
+    for key in partial_config:
+        if key in config and isinstance(partial_config[key], dict) and isinstance(config[key], dict):
+            partial_update_config(config[key], partial_config[key])
+        else:
+            config[key] = partial_config[key]
+    return config
+
+
+def resolve_and_load_config(path: str, config_name="config.yaml") -> dict:
+    path = os.path.realpath(os.path.expanduser(path))
+    if os.path.isdir(path):
+        config_path = os.path.join(path, config_name)
+    else:
+        config_path = path
+    if os.path.isfile(config_path):
+        pass
+    else:
+        raise Exception(f"Cannot find a valid config at {path}")
+    config = load_config(config_path)
+    return config
+
+
+class SafeLoaderWithTuple(yaml.SafeLoader):
+    """A yaml safe loader with python tuple loading capabilities."""
+
+    def construct_python_tuple(self, node):
+        return tuple(self.construct_sequence(node))
+
+
+SafeLoaderWithTuple.add_constructor("tag:yaml.org,2002:python/tuple", SafeLoaderWithTuple.construct_python_tuple)
+
+
+def load_config(filename: str) -> dict:
+    """Load a yaml file."""
+    filename = os.path.realpath(os.path.expanduser(filename))
+    return yaml.load(open(filename), Loader=SafeLoaderWithTuple)
+
+
+def dump_config(config: dict, filename: str) -> None:
+    """Dump a config file"""
+    filename = os.path.realpath(os.path.expanduser(filename))
+    yaml.dump(config, open(filename, "w"), sort_keys=False)

segment_everything/vendored/efficientvit/apps/utils/opt.py

@@ -0,0 +1,28 @@
+# EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
+# Han Cai, Junyan Li, Muyan Hu, Chuang Gan, Song Han
+# International Conference on Computer Vision (ICCV), 2023
+
+import torch
+
+__all__ = ["REGISTERED_OPTIMIZER_DICT", "build_optimizer"]
+
+# register optimizer here
+#   name: optimizer, kwargs with default values
+REGISTERED_OPTIMIZER_DICT: dict[str, tuple[type, dict[str, any]]] = {
+    "sgd": (torch.optim.SGD, {"momentum": 0.9, "nesterov": True}),
+    "adam": (torch.optim.Adam, {"betas": (0.9, 0.999), "eps": 1e-8, "amsgrad": False}),
+    "adamw": (torch.optim.AdamW, {"betas": (0.9, 0.999), "eps": 1e-8, "amsgrad": False}),
+}
+
+
+def build_optimizer(
+    net_params, optimizer_name: str, optimizer_params: dict or None, init_lr: float
+) -> torch.optim.Optimizer:
+    optimizer_class, default_params = REGISTERED_OPTIMIZER_DICT[optimizer_name]
+    optimizer_params = optimizer_params or {}
+
+    for key in default_params:
+        if key in optimizer_params:
+            default_params[key] = optimizer_params[key]
+    optimizer = optimizer_class(net_params, init_lr, **default_params)
+    return optimizer

segment_everything/vendored/efficientvit/cls_model_zoo.py

@@ -0,0 +1,79 @@
+# EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
+# Han Cai, Junyan Li, Muyan Hu, Chuang Gan, Song Han
+# International Conference on Computer Vision (ICCV), 2023
+
+from efficientvit.models.efficientvit import (
+    EfficientViTCls,
+    efficientvit_cls_b0,
+    efficientvit_cls_b1,
+    efficientvit_cls_b2,
+    efficientvit_cls_b3,
+    efficientvit_cls_l1,
+    efficientvit_cls_l2,
+    efficientvit_cls_l3,
+)
+from efficientvit.models.nn.norm import set_norm_eps
+from efficientvit.models.utils import load_state_dict_from_file
+
+__all__ = ["create_cls_model"]
+
+
+REGISTERED_CLS_MODEL: dict[str, str] = {
+    "b0-r224": "assets/checkpoints/cls/b0-r224.pt",
+    ###############################################
+    "b1-r224": "assets/checkpoints/cls/b1-r224.pt",
+    "b1-r256": "assets/checkpoints/cls/b1-r256.pt",
+    "b1-r288": "assets/checkpoints/cls/b1-r288.pt",
+    ###############################################
+    "b2-r224": "assets/checkpoints/cls/b2-r224.pt",
+    "b2-r256": "assets/checkpoints/cls/b2-r256.pt",
+    "b2-r288": "assets/checkpoints/cls/b2-r288.pt",
+    ###############################################
+    "b3-r224": "assets/checkpoints/cls/b3-r224.pt",
+    "b3-r256": "assets/checkpoints/cls/b3-r256.pt",
+    "b3-r288": "assets/checkpoints/cls/b3-r288.pt",
+    ###############################################
+    "l1-r224": "assets/checkpoints/cls/l1-r224.pt",
+    ###############################################
+    "l2-r224": "assets/checkpoints/cls/l2-r224.pt",
+    "l2-r256": "assets/checkpoints/cls/l2-r256.pt",
+    "l2-r288": "assets/checkpoints/cls/l2-r288.pt",
+    "l2-r320": "assets/checkpoints/cls/l2-r320.pt",
+    "l2-r384": "assets/checkpoints/cls/l2-r384.pt",
+    ###############################################
+    "l3-r224": "assets/checkpoints/cls/l3-r224.pt",
+    "l3-r256": "assets/checkpoints/cls/l3-r256.pt",
+    "l3-r288": "assets/checkpoints/cls/l3-r288.pt",
+    "l3-r320": "assets/checkpoints/cls/l3-r320.pt",
+    "l3-r384": "assets/checkpoints/cls/l3-r384.pt",
+}
+
+
+def create_cls_model(name: str, pretrained=True, weight_url: str or None = None, **kwargs) -> EfficientViTCls:
+    model_dict = {
+        "b0": efficientvit_cls_b0,
+        "b1": efficientvit_cls_b1,
+        "b2": efficientvit_cls_b2,
+        "b3": efficientvit_cls_b3,
+        #########################
+        "l1": efficientvit_cls_l1,
+        "l2": efficientvit_cls_l2,
+        "l3": efficientvit_cls_l3,
+    }
+
+    model_id = name.split("-")[0]
+    if model_id not in model_dict:
+        raise ValueError(f"Do not find {name} in the model zoo. List of models: {list(model_dict.keys())}")
+    else:
+        model = model_dict[model_id](**kwargs)
+    if model_id in ["l1", "l2", "l3"]:
+        set_norm_eps(model, 1e-7)
+
+    if pretrained:
+        weight_url = weight_url or REGISTERED_CLS_MODEL.get(name, None)
+        if weight_url is None:
+            raise ValueError(f"Do not find the pretrained weight of {name}.")
+        else:
+            weight = load_state_dict_from_file(weight_url)
+            model.load_state_dict(weight)
+    return model

segment_everything/vendored/efficientvit/clscore/data_provider/__init__.py

@@ -0,0 +1,5 @@
+# EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
+# Han Cai, Junyan Li, Muyan Hu, Chuang Gan, Song Han
+# International Conference on Computer Vision (ICCV), 2023
+
+from .imagenet import *

segment_everything/vendored/efficientvit/clscore/data_provider/imagenet.py

@@ -0,0 +1,142 @@
+# EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
+# Han Cai, Junyan Li, Muyan Hu, Chuang Gan, Song Han
+# International Conference on Computer Vision (ICCV), 2023
+
+import copy
+import math
+import os
+
+import torchvision.transforms as transforms
+from torchvision.datasets import ImageFolder
+
+from ...apps.data_provider import DataProvider
+from ...apps.data_provider.augment import RandAug
+from ...apps.data_provider.random_resolution import (
+    MyRandomResizedCrop,
+    get_interpolate,
+)
+from ...apps.utils import partial_update_config
+from ...models.utils import val2list
+
+__all__ = ["ImageNetDataProvider"]
+
+
+class ImageNetDataProvider(DataProvider):
+    name = "imagenet"
+
+    data_dir = "/dataset/imagenet"
+    n_classes = 1000
+    _DEFAULT_RRC_CONFIG = {
+        "train_interpolate": "random",
+        "test_interpolate": "bicubic",
+        "test_crop_ratio": 1.0,
+    }
+
+    def __init__(
+        self,
+        data_dir: str or None = None,
+        rrc_config: dict or None = None,
+        data_aug: dict or list[dict] or None = None,
+        ###########################################
+        train_batch_size=128,
+        test_batch_size=128,
+        valid_size: int or float or None = None,
+        n_worker=8,
+        image_size: int or list[int] = 224,
+        num_replicas: int or None = None,
+        rank: int or None = None,
+        train_ratio: float or None = None,
+        drop_last: bool = False,
+    ):
+        self.data_dir = data_dir or self.data_dir
+        self.rrc_config = partial_update_config(
+            copy.deepcopy(self._DEFAULT_RRC_CONFIG),
+            rrc_config or {},
+        )
+        self.data_aug = data_aug
+
+        super().__init__(
+            train_batch_size,
+            test_batch_size,
+            valid_size,
+            n_worker,
+            image_size,
+            num_replicas,
+            rank,
+            train_ratio,
+            drop_last,
+        )
+
+    def build_valid_transform(
+        self, image_size: tuple[int, int] or None = None
+    ) -> any:
+        image_size = (image_size or self.active_image_size)[0]
+        crop_size = int(
+            math.ceil(image_size / self.rrc_config["test_crop_ratio"])
+        )
+        return transforms.Compose(
+            [
+                transforms.Resize(
+                    crop_size,
+                    interpolation=get_interpolate(
+                        self.rrc_config["test_interpolate"]
+                    ),
+                ),
+                transforms.CenterCrop(image_size),
+                transforms.ToTensor(),
+                transforms.Normalize(**self.mean_std),
+            ]
+        )
+
+    def build_train_transform(
+        self, image_size: tuple[int, int] or None = None
+    ) -> any:
+        image_size = image_size or self.image_size
+
+        # random_resize_crop -> random_horizontal_flip
+        train_transforms = [
+            MyRandomResizedCrop(
+                interpolation=self.rrc_config["train_interpolate"]
+            ),
+            transforms.RandomHorizontalFlip(),
+        ]
+
+        # data augmentation
+        post_aug = []
+        if self.data_aug is not None:
+            for aug_op in val2list(self.data_aug):
+                if aug_op["name"] == "randaug":
+                    data_aug = RandAug(aug_op, mean=self.mean_std["mean"])
+                elif aug_op["name"] == "erase":
+                    from timm.data.random_erasing import RandomErasing
+
+                    random_erase = RandomErasing(aug_op["p"], device="cpu")
+                    post_aug.append(random_erase)
+                    data_aug = None
+                else:
+                    raise NotImplementedError
+                if data_aug is not None:
+                    train_transforms.append(data_aug)
+        train_transforms = [
+            *train_transforms,
+            transforms.ToTensor(),
+            transforms.Normalize(**self.mean_std),
+            *post_aug,
+        ]
+        return transforms.Compose(train_transforms)
+
+    def build_datasets(self) -> tuple[any, any, any]:
+        train_transform = self.build_train_transform()
+        valid_transform = self.build_valid_transform()
+
+        train_dataset = ImageFolder(
+            os.path.join(self.data_dir, "train"), train_transform
+        )
+        test_dataset = ImageFolder(
+            os.path.join(self.data_dir, "val"), valid_transform
+        )
+
+        train_dataset, val_dataset = self.sample_val_dataset(
+            train_dataset, valid_transform
+        )
+        return train_dataset, val_dataset, test_dataset

segment_everything/vendored/efficientvit/clscore/trainer/__init__.py

@@ -0,0 +1,6 @@
+# EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
+# Han Cai, Junyan Li, Muyan Hu, Chuang Gan, Song Han
+# International Conference on Computer Vision (ICCV), 2023
+
+from .cls_run_config import *
+from .cls_trainer import *

segment_everything/vendored/efficientvit/clscore/trainer/cls_run_config.py

@@ -0,0 +1,18 @@
+# EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
+# Han Cai, Junyan Li, Muyan Hu, Chuang Gan, Song Han
+# International Conference on Computer Vision (ICCV), 2023
+
+from ...apps.trainer.run_config import RunConfig
+
+__all__ = ["ClsRunConfig"]
+
+
+class ClsRunConfig(RunConfig):
+    label_smooth: float
+    mixup_config: dict  # allow none to turn off mixup
+    bce: bool
+    mesa: dict
+
+    @property
+    def none_allowed(self):
+        return ["mixup_config", "mesa"] + super().none_allowed