Fix: canonicalization

sources/unipercept/cli/run.py

@@ -0,0 +1,132 @@
+"""
+Run a model in realtime or on a previously saved directory of images.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import typing as T
+
+import torch
+from omegaconf import DictConfig
+from tabulate import tabulate
+
+import unipercept as up
+from unipercept.cli._command import command
+
+_logger = up.log.get_logger()
+
+
+KEY_SESSION_ID = "session_id"
+
+
+@command(help="trian a model", description=__doc__)
+@command.with_config
+def train(p: argparse.ArgumentParser):
+    p_size = p.add_mutually_exclusive_group(required=False)
+    p_size.add_argument(
+        "--size", type=int, help="Size of the input images in pixels (smallest side)"
+    )
+    p.add_argument(
+        "--weights",
+        "-w",
+        type=str,
+        help="path to load model weights from (overrides any state recovered by the config)",
+    )
+    p.add_argument(
+        "--render",
+        type=str,
+        default="segmentation",
+        choices=["segmentation", "depth", "noop"],
+        help="rendering mode",
+    )
+
+    p.add_argument("input", type=str, default="0", help="input stream or directory")
+
+    return _main
+
+
+@torch.inference_mode()
+def _main(args: argparse.Namespace):
+    config = args.config
+
+    model = up.models.load(config.model)
+    preprocess = _build_transforms()
+
+    if up.file_io.isdir(args.input):
+        run = _run_filesystem(model, preprocess, args.input)
+    else:
+        cap_num = int(args.input)
+        cap = _get_capture(cap_num)
+        run = _run_realtime(model, preprocess, cap)
+
+    for inp, out in run:
+        print(out)
+
+
+def _build_transforms(args):
+    import torchvision.transforms.v2 as transforms
+
+    tf = []
+    if args.size:
+        tf.append(transforms.Resize(args.size))
+
+    return up.data.ops.TorchvisionOp(tf)
+
+
+def _get_capture(cap_num):
+    import cv2
+
+    cap = cv2.VideoCapture(cap_num, cv2.CAP_V4L2)
+    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 224)
+    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 224)
+    cap.set(cv2.CAP_PROP_FPS, 1)
+    return cap
+
+
+def _run_realtime(model, preprocess, cap):
+    import numpy as np
+    import torchvision.transforms.v2 as transforms
+
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+    frame_num = 0
+    while True:
+        ret, img_np = cap.read()
+        if not ret:
+            break
+
+        # BGR -> RGB
+        img = transforms.functional.to_tensor(img_np[..., [2, 1, 0]])
+        inp = up.create_inputs(img, frame_offset=frame_num)
+        inp = preprocess(inp)
+        out = model(inp)
+
+        yield inp, out
+
+        frame_num += 1
+
+
+def _run_filesystem(model, preprocess, path):
+    root = up.file_io.Path(path)
+    root_paths = list(root.iterdir())
+
+    if all(p.is_dir() for p in root_paths):
+        for p in root_paths:
+            yield from _run_filesystem(model, preprocess, p)
+    elif all(p.name.endswith(".png") for p in root_paths):
+        for p in root_paths:
+            img = up.data.tensors.Image.read(p)
+            inp = up.create_inputs(img, frame_offset=0)
+            inp = preprocess(inp)
+            out = model(inp)
+
+            yield inp, out
+    else:
+        msg = (
+            f"Invalid directory structure: {str(path)!r}, expected directories (sequence) "
+            "of PNG images sortable by name!"
+        )
+        raise ValueError(msg)

sources/unipercept/data/_loader.py

@@ -3,13 +3,11 @@
 from __future__ import annotations
 
 import abc
-import dataclasses
 import dataclasses as D
 import enum
 import functools
 import itertools
 import math
-import multiprocessing as M
 import operator
 import typing as T
 import warnings
@@ -47,41 +45,34 @@
 _logger = get_logger(__name__)
 
 
-def _suggest_workers():
-    """
-    Suggests the number of workers for the dataloader based on the number of available CPUs
-    """
-    try:
-        return max(cpus_available() - (cpus_available() // 4), 1)
-    except Exception:
-        return M.cpu_count() // get_process_count()
-
-
-DEFAULT_NUM_WORKERS = get_env(
-    int,
-    "UP_DATALOADER_WORKERS",
-    "UNIPERCEPT_DATALOADER_WORKERS",
-    default=_suggest_workers(),
-)
-
-DEFAULT_PREFETCH_FACTOR = get_env(
-    int,
-    "UP_DATALOADER_PREFETCH_FACTOR",
-    "UNIPERCEPT_DATALOADER_PREFETCH_FACTOR",
-    default=2,
-)
-
-
-@dataclasses.dataclass(slots=True, frozen=True)
+@D.dataclass(slots=True, frozen=True)
 class DataLoaderConfig:
     """
     Configuration parameters passed to the PyTorch dataoader
     """
 
     drop_last: bool = False
-    pin_memory: bool = True
-    num_workers: int = DEFAULT_NUM_WORKERS
-    prefetch_factor: int | None = DEFAULT_PREFETCH_FACTOR
+    pin_memory: bool = D.field(
+        default_factory=lambda: get_env(
+            bool,
+            "UP_DATALOADER_PIN_MEMORY",
+            default=False,
+        )
+    )
+    num_workers: int = D.field(
+        default_factory=lambda: get_env(
+            int,
+            "UP_DATALOADER_WORKERS",
+            default=max(cpus_available(), 16),
+        )
+    )
+    prefetch_factor: int | None = D.field(
+        default_factory=lambda: get_env(
+            int,
+            "UP_DATALOADER_PREFETCH_FACTOR",
+            default=2,
+        )
+    )
     persistent_workers: bool | None = False
 
 
@@ -90,7 +81,7 @@ class DataLoaderConfig:
 ##################
 
 
-@dataclasses.dataclass(slots=True, frozen=True)
+@D.dataclass(slots=True, frozen=True)
 class DataLoaderFactory:
     """
     Factory for creating dataloaders.
@@ -115,8 +106,8 @@ class DataLoaderFactory:
     dataset: PerceptionDataset
     actions: T.Sequence[Op]
     sampler: SamplerFactory
-    config: DataLoaderConfig = dataclasses.field(default_factory=DataLoaderConfig)
-    iterable: bool = dataclasses.field(
+    config: DataLoaderConfig = D.field(default_factory=DataLoaderConfig)
+    iterable: bool = D.field(
         default=False,
         metadata={
             "help": (
@@ -142,11 +133,10 @@ def __call__(
 
         # Keyword arguments for the loader
         loader_kwargs = {
-            k: v for k, v in dataclasses.asdict(self.config).items() if v is not None
+            k: v for k, v in D.asdict(self.config).items() if v is not None
         }
 
         # Instantiate sampler
-
         sampler_kwargs = {}
         if not use_distributed:
             sampler_kwargs["process_count"] = 1
@@ -187,7 +177,12 @@ def __call__(
         # Loader
         loader_kwargs["batch_size"] = batch_size
         loader_kwargs.setdefault("collate_fn", InputData.collate)
-        loader_kwargs.setdefault("worker_init_fn", _worker_init_fn)
+
+        if loader_kwargs["num_workers"] > 0:
+            loader_kwargs.setdefault("worker_init_fn", _worker_init_fn)
+        elif "worker_init_fn" in loader_kwargs:
+            msg = f"Worker init function is set, but num_workers is {loader_kwargs['num_workers']}."
+            raise ValueError(msg)
 
         _logger.debug(
             "Creating dataloader (%d queued; %d × %d items):\n%s",
@@ -683,7 +678,3 @@ def __init__(
 
     def __call__(self, queue: PerceptionDataqueue) -> Sampler:
         return self._fn(queue)
-
-
-if __name__ == "__main__":
-    print("Default configuration for dataloader workers: ", DEFAULT_NUM_WORKERS)

sources/unipercept/integrations/wandb_integration.py

@@ -368,7 +368,10 @@ def artifact_historic_delete(artifact: wandb.Artifact, keep: int) -> None:
 
     api = wandb.Api()
 
-    vs = api.artifact_versions(type_name=artifact.type, name=name)
+    vs = [a.version for a in api.artifacts(type_name=artifact.type, name=name)]
+    if len(vs) <= keep:
+        _logger.info(f"Keeping all {name} artifacts")
+        return
     vs = sorted(vs, key=artifact_version_as_int, reverse=True)
     for artifact in vs[keep:]:
         _logger.info(f"Deleting artifact {name} version {artifact.version}")

sources/unipercept/model.py

@@ -346,7 +346,7 @@ def __call__(
                 overrides_list = list(overrides)
             _logger.info(
                 "Instantiating model with configuration overrides: %s",
-                ", ".join(overrides_list),
+                ", ".join(overrides_list) if len(overrides_list) > 0 else "(none)",
             )
             model_config = apply_overrides(model_config, overrides_list)
         else:

sources/unipercept/nn/layers/conv/_conditional.py

@@ -17,7 +17,7 @@
 
 from unipercept.utils.function import to_2tuple
 
-from .utils import NormActivationMixin, PaddingMixin
+from .utils import NormActivationMixin
 
 
 def get_condconv_initializer(initializer, num_experts, expert_shape):
@@ -40,7 +40,7 @@ def condconv_initializer(weight):
     return condconv_initializer
 
 
-class CondConv2d(PaddingMixin, NormActivationMixin, nn.Module):
+class CondConv2d(NormActivationMixin, nn.Module):
     r"""
     Based on the implementation in `timm.layers`, where their docs state:
     >   Conditionally Parameterized Convolution

sources/unipercept/nn/layers/conv/_deform.py

@@ -20,7 +20,7 @@
 
 from ..weight import init_xavier_fill_
 from ._extended import Conv2d
-from .utils import NormActivationMixin, PaddingMixin
+from .utils import NormActivationMixin
 
 __all__ = ["ModDeform2d", "DeformConv2d"]
 
@@ -223,7 +223,7 @@ def __repr__(self) -> str:
         return s
 
 
-class ModDeform2d(NormActivationMixin, PaddingMixin, nn.Module):
+class ModDeform2d(NormActivationMixin, nn.Module):
     """
     Modulated deformable convolution.
     The offset mask is computed by  a convolutional layer.

sources/unipercept/nn/layers/conv/_extended.py

@@ -7,23 +7,15 @@
 import torch.nn as nn
 import typing_extensions as TX
 
-from .utils import NormActivationMixin, PaddingMixin
+from .utils import NormActivationMixin
 
-__all__ = ["Conv2d", "Standard2d", "PadConv2d", "Separable2d"]
+__all__ = ["Conv2d", "Standard2d", "Separable2d"]
 
 
 class Conv2d(NormActivationMixin, nn.Conv2d):
     pass
 
 
-class PadConv2d(PaddingMixin, Conv2d):
-    @TX.override
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self._padding_forward(x, self.kernel_size, self.stride, self.dilation)
-        x = self._conv_forward(x, self.weight, self.bias)
-        return x
-
-
 class Standard2d(Conv2d):
     """
     Implements weight standardization with learnable gain.
@@ -43,7 +35,7 @@ def __init__(self, *args, gamma=1.0, eps=1e-6, gain=1.0, **kwargs):
 
     @TX.override
     def forward(self, x):
-        weight = F.batch_norm(
+        weight = nn.functional.batch_norm(
             self.weight.reshape(1, self.out_channels, -1),
             None,
             None,