easydetect 0.2.1__tar.gz → 0.2.2__tar.gz

{easydetect-0.2.1 → easydetect-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: easydetect
-Version: 0.2.1
+Version: 0.2.2
 Summary: Easy real-time object detection you can ship: D-FINE with PyTorch training, OpenVINO inference on CPU/GPU/NPU, and a browser app to label and train. Apache-2.0 throughout.
 Project-URL: Homepage, https://github.com/themakerrobot/easydetect
 Project-URL: Source, https://github.com/themakerrobot/easydetect

{easydetect-0.2.1 → easydetect-0.2.2}/easydetect/__init__.py

@@ -24,7 +24,7 @@ needs numpy/opencv/openvino/pyyaml; training adds
 
 from __future__ import annotations
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"
 
 from .errors import DownloadError, EasyDetectError, ModelNotFoundError
 from .metrics import BoxMetrics, DetMetrics

{easydetect-0.2.1 → easydetect-0.2.2}/easydetect/exporter.py

@@ -62,12 +62,30 @@ def export_onnx(net, names, imgsz=640, out_dir=".", fname="easydetect", half=Fal
             opset_version=17,
             dynamo=False,
         )
-    _write_labels(out_dir, fname, names)
+    table = _write_labels(out_dir, fname, names)
+    _embed_names(onnx_path, table)
     if verbose:
         print(f"[easydetect] exported: {onnx_path}")
     return onnx_path
 
 
+def _embed_names(onnx_path: Path, table: dict[int, str]) -> None:
+    """Keep the class names inside the .onnx too, so the file works on its own —
+    one download from a Hugging Face page, no labels.txt to fetch beside it."""
+    import onnx
+
+    from .predictor import ONNX_NAMES_KEY
+
+    model = onnx.load(str(onnx_path))
+    kept = [p for p in model.metadata_props if p.key != ONNX_NAMES_KEY]
+    del model.metadata_props[:]
+    model.metadata_props.extend(kept)
+    entry = model.metadata_props.add()
+    entry.key = ONNX_NAMES_KEY
+    entry.value = json.dumps({str(k): v for k, v in table.items()}, ensure_ascii=False)
+    onnx.save(model, str(onnx_path))
+
+
 def export_openvino(net, names, imgsz=640, out_dir=".", fname="easydetect", half=False,
                     verbose=True):
     """Write ``<out_dir>/<fname>.xml`` (+ .bin, + labels). Returns the .xml path."""

{easydetect-0.2.1 → easydetect-0.2.2}/easydetect/model.py

@@ -349,7 +349,7 @@ class Detector:
         name: str = "train",
         resume: bool = False,
         patience: int = 50,
-        lr0: float = 1e-4,
+        lr0: float | None = None,
         seed: int = 0,
         val: bool = True,
         amp: bool = True,
@@ -362,7 +362,9 @@ class Detector:
         """Train on a data.yaml dataset. Returns the path of ``best.pt``.
 
         ``freeze="backbone"`` trains about twice as fast on CPU and helps when
-        the dataset is small. ``augment=True`` trains on zoomed-out, cropped
+        the dataset is small. ``lr0=None`` sets the learning rate from the
+        batch size, ``1e-4 × √(batch / 4)``, so a bigger batch is not a run
+        that learned less; a number sets it outright. ``augment=True`` trains on zoomed-out, cropped
         and colour-jittered pictures (D-FINE's recipe) and on plain ones for
         the last tenth of the epochs; ``False`` keeps only the flip.
         ``on_epoch_end`` is called with a dict of the epoch's numbers — the

{easydetect-0.2.1 → easydetect-0.2.2}/easydetect/predictor.py

@@ -37,13 +37,87 @@ def _looks_like_logits(scores: np.ndarray) -> bool:
     return bool(scores.size) and (scores.min() < 0.0 or scores.max() > 1.0)
 
 
+#: The ONNX metadata key easydetect keeps the class names under.
+ONNX_NAMES_KEY = "easydetect.names"
+
+
+def _varint(f) -> int | None:
+    shift = value = 0
+    while True:
+        byte = f.read(1)
+        if not byte:
+            return None
+        value |= (byte[0] & 0x7F) << shift
+        if byte[0] < 0x80:
+            return value
+        shift += 7
+
+
+def onnx_metadata(path: Path) -> dict[str, str]:
+    """The ``metadata_props`` of an ONNX file, read without the ``onnx`` package.
+
+    ModelProto keeps them as field 14, each a (key = 1, value = 2) pair. Only
+    the top level of the file is walked; the graph and its weights are skipped
+    over with a seek, so a large model costs no more than a small one.
+    """
+    props: dict[str, str] = {}
+    with open(path, "rb") as f:
+        while (tag := _varint(f)) is not None:
+            field, wire = tag >> 3, tag & 7
+            if wire == 0:
+                _varint(f)
+            elif wire == 1:
+                f.seek(8, 1)
+            elif wire == 5:
+                f.seek(4, 1)
+            elif wire == 2:
+                size = _varint(f) or 0
+                if field != 14:
+                    f.seek(size, 1)
+                    continue
+                entry, key, value, i = f.read(size), "", "", 0
+                while i < len(entry):
+                    sub_tag, n = entry[i], 0
+                    i += 1
+                    shift = 0
+                    while True:                       # the length, a varint
+                        n |= (entry[i] & 0x7F) << shift
+                        i += 1
+                        if entry[i - 1] < 0x80:
+                            break
+                        shift += 7
+                    text = entry[i:i + n].decode("utf-8", "replace")
+                    i += n
+                    if sub_tag >> 3 == 1:
+                        key = text
+                    elif sub_tag >> 3 == 2:
+                        value = text
+                props[key] = value
+            else:                                     # not an ONNX file after all
+                break
+    return props
+
+
 def read_names(model_path: Path) -> dict[int, str]:
-    """Class names from ``labels.txt`` or ``<stem>.names.json`` beside the model."""
+    """Class names: ``<stem>.names.json`` beside the model, then the names an
+    ``.onnx`` carries inside it, then ``labels.txt`` in its folder.
+
+    The names inside the file come before a ``labels.txt`` because they belong
+    to that file: a lone ``best.onnx`` downloaded into a folder that holds
+    another model's ``labels.txt`` still names its own classes.
+    """
     model_path = Path(model_path)
     sidecar = model_path.with_suffix(".names.json")
     if sidecar.exists():
         table = json.loads(sidecar.read_text(encoding="utf-8"))
         return {int(k): str(v) for k, v in table.items()}
+    if model_path.suffix.lower() == ".onnx" and model_path.is_file():
+        try:
+            inside = onnx_metadata(model_path).get(ONNX_NAMES_KEY)
+            if inside:
+                return {int(k): str(v) for k, v in json.loads(inside).items()}
+        except (OSError, ValueError, IndexError):
+            pass
     labels = model_path.parent / "labels.txt"
     if labels.exists():
         lines = [ln.strip() for ln in labels.read_text(encoding="utf-8").splitlines()]

{easydetect-0.2.1 → easydetect-0.2.2}/easydetect/trainer.py

@@ -1,7 +1,8 @@
 # Apache-2.0
 """Training loop (PyTorch). Kept deliberately small and readable.
 
-AdamW with a lower LR on the backbone, linear warmup into a cosine decay, AMP
+AdamW with a lower LR on the backbone, the LR set from the batch size unless
+given, linear warmup into a cosine decay, AMP
 on CUDA, ``last.pt``/``best.pt`` after every epoch, early stop on ``patience``,
 and ``resume=True`` to pick a killed run back up where it stopped.
 """
@@ -31,6 +32,24 @@ def seed_everything(seed: int) -> None:
     torch.cuda.manual_seed_all(seed)
 
 
+#: The learning rate a batch of 4 fine-tunes well at: D-FINE from its COCO
+#: weights reached 0.765 mAP50-95 on a dataset where RT-DETR's own run gave 0.762.
+REF_BATCH, REF_LR = 4, 1e-4
+
+
+def auto_lr(batch: int) -> float:
+    """The learning rate for a batch size: ``1e-4 × √(batch / 4)``.
+
+    A bigger batch means fewer optimizer steps an epoch, and Adam moves each
+    weight by about the learning rate per step whatever the batch, so at one
+    fixed rate a batch of 32 ends the same epochs having learned an eighth as
+    much. The square root is the usual rule for Adam (four times the batch,
+    twice the rate); at 32 it gives 2.8e-4, beside the 2–2.5e-4 D-FINE's own
+    configs train with at their batch of 32.
+    """
+    return REF_LR * math.sqrt(max(int(batch), 1) / REF_BATCH)
+
+
 #: The backbone learns slower than the rest — D-FINE's per-size ratios.
 BACKBONE_LR_MULT = {"n": 0.5, "s": 0.5, "m": 0.1, "l": 0.05, "x": 0.01}
 
@@ -153,7 +172,7 @@ class Trainer:
         epochs=100,
         imgsz=640,
         batch=8,
-        lr=1e-4,
+        lr=None,
         lr_backbone_mult=None,
         weight_decay=1e-4,
         warmup_epochs=1,
@@ -174,6 +193,8 @@ class Trainer:
         seed_everything(seed)
         self.augment = augment
         self.seed, self.freeze, self.origin = seed, freeze, origin
+        self.lr_auto = lr is None
+        lr = auto_lr(batch) if lr is None else float(lr)
         self.lr, self.lr_backbone_mult, self.weight_decay = lr, lr_backbone_mult, weight_decay
         self.net = net
         self.frozen = freeze_modules(net, freeze)
@@ -288,6 +309,8 @@ class Trainer:
             "optimizer": {
                 "name": "AdamW",
                 "lr": self.lr,
+                "lr_from": (f"batch {self.batch}: {REF_LR:g} × √({self.batch}/{REF_BATCH})"
+                            if self.lr_auto else "set"),
                 "lr_backbone": self.lr * self.lr_backbone_mult,
                 "weight_decay": self.weight_decay,
                 "schedule": "linear warmup, then cosine decay to 1%",

{easydetect-0.2.1 → easydetect-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "easydetect"
-version = "0.2.1"
+version = "0.2.2"
 description = "Easy real-time object detection you can ship: D-FINE with PyTorch training, OpenVINO inference on CPU/GPU/NPU, and a browser app to label and train. Apache-2.0 throughout."
 readme = "README.md"
 requires-python = ">=3.10"

{easydetect-0.2.1 → easydetect-0.2.2}/tests/test_onnxruntime.py

@@ -104,3 +104,30 @@ def test_a_named_model_downloads_the_onnx_for_onnxruntime(monkeypatch, tmp_path,
     monkeypatch.setattr(downloads, "_asset", fake_asset)
     assert downloads.download_onnx("D-FINE-S") == tmp_path / "dfine-s.onnx"
     assert fetched == ["dfine-s.onnx", "labels.txt"]
+
+
+def test_an_onnx_carries_its_class_names_so_it_works_alone(tiny_ir, tmp_path):
+    """One file from a Hugging Face page: no labels.txt beside it, or someone else's."""
+    import shutil
+
+    lone = tmp_path / "downloads" / "best.onnx"
+    lone.parent.mkdir()
+    shutil.copy(tiny_ir.with_suffix(".onnx"), lone)
+    (lone.parent / "labels.txt").write_text("person\ncar\nbus\n")   # another model's
+    for backend in ("openvino", "onnxruntime"):
+        model = Detector(str(lone), backend=backend, verbose=False)
+        assert model(draw(), conf=0.0)[0].names == {0: "can", 1: "bottle"}, backend
+
+
+def test_the_names_are_read_without_the_onnx_package(tiny_ir, tmp_path):
+    onnx = pytest.importorskip("onnx")
+    from easydetect.predictor import ONNX_NAMES_KEY, onnx_metadata, read_names
+
+    path = tiny_ir.with_suffix(".onnx")
+    expected = {p.key: p.value for p in onnx.load(str(path)).metadata_props}
+    assert onnx_metadata(path) == expected and ONNX_NAMES_KEY in expected
+
+    # not an ONNX file at all: no names, and no crash
+    junk = tmp_path / "junk.onnx"
+    junk.write_bytes(b"\xff\xff\xff\xff\x0f not a model")
+    assert read_names(junk) == {}

{easydetect-0.2.1 → easydetect-0.2.2}/tests/test_train_export.py

@@ -217,6 +217,9 @@ def test_training_reports_every_epoch_to_a_callback_and_a_csv(dataset, tmp_path,
     assert setup["variant"] == "n" and setup["start"] == {"kind": "scratch"}
     assert setup["freeze"] == "backbone" and 0 < setup["trainable_params"] < setup["params"]
     assert setup["optimizer"]["name"] == "AdamW" and setup["amp"] is False
+    # no lr0: the rate came from the batch of 2, and the record says so
+    assert setup["optimizer"]["lr"] == pytest.approx(1e-4 * (2 / 4) ** 0.5)
+    assert setup["optimizer"]["lr_from"].startswith("batch 2:")
     train = setup["data"]["train"]
     assert train["images"] > 0 and sum(train["per_class"]) == train["boxes"] > 0
     assert setup["versions"]["torch"] and setup["device"] == "cpu" and setup["gpu"] is None
@@ -306,3 +309,27 @@ def test_pretrained_false_never_reaches_for_the_mirror(dataset, tmp_path, monkey
         data=str(dataset), project=str(tmp_path), epochs=1, imgsz=64, batch=2, workers=0,
         device="cpu", val=False, amp=False,
     ).exists()
+
+
+def test_the_learning_rate_follows_the_batch_unless_it_is_given():
+    from easydetect.trainer import auto_lr
+
+    assert auto_lr(4) == pytest.approx(1e-4)
+    assert auto_lr(16) == pytest.approx(2e-4)
+    assert auto_lr(32) == pytest.approx(2.83e-4, rel=1e-3)
+    assert auto_lr(0) == auto_lr(1)
+
+
+@needs_torch
+def test_a_learning_rate_that_is_given_is_used_as_it_is(dataset, tmp_path, monkeypatch):
+    import json
+
+    monkeypatch.setattr("easydetect.nn.hgnetv2.HGNetv2.load_imagenet", lambda self, name, url: None)
+    model = Detector("dfine-n", verbose=False, pretrained=False)
+    best = model.train(data=str(dataset), project=str(tmp_path), epochs=1, imgsz=64, batch=2,
+                       workers=0, device="cpu", val=False, amp=False, lr0=3e-4,
+                       warmup_epochs=0, weight_decay=0.0, lr_backbone_mult=0.25)
+    opt = json.loads((best.parent.parent / "run.json").read_text())["optimizer"]
+    assert opt["lr"] == pytest.approx(3e-4) and opt["lr_from"] == "set"
+    assert opt["lr_backbone"] == pytest.approx(0.75e-4)
+    assert opt["warmup_epochs"] == 0 and opt["weight_decay"] == 0.0