easydetect 0.2.1__tar.gz → 0.2.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: easydetect
-Version: 0.2.1
+Version: 0.2.3
 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.3}/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.3"
 
 from .errors import DownloadError, EasyDetectError, ModelNotFoundError
 from .metrics import BoxMetrics, DetMetrics

{easydetect-0.2.1 → easydetect-0.2.3}/easydetect/cli.py

@@ -29,7 +29,7 @@ Modes:
 Common keys:
   model=dfine-n|s|m|l|x|best.pt|model.xml   source=bus.jpg|dir|video.mp4|0|url
   data=data.yaml  epochs=100  imgsz=640  batch=8  conf=0.25  iou=0.7|none  device=0|cpu|AUTO
-  project=runs  name=predict  save=true  show=false  format=openvino  half=true
+  contain=0.8  project=runs  name=predict  save=true  show=false  format=openvino  half=true
 
 Examples:
   easydetect predict model=dfine-s source=bus.jpg conf=0.5

{easydetect-0.2.1 → easydetect-0.2.3}/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.3}/easydetect/model.py

@@ -120,6 +120,7 @@ class Detector:
         max_det: int = 300,
         classes: list[int] | None = None,
         iou: float | None = 0.7,
+        contain: float | None = None,
         stream: bool = False,
         save: bool = False,
         show: bool = False,
@@ -136,6 +137,9 @@ class Detector:
         only sane option for a long video or a live camera). ``iou`` drops a
         box covering a higher-scoring one by more than that IoU — the same
         object found twice, sometimes under two classes; ``None`` keeps all.
+        ``contain=0.8`` also merges a box with 80% of its area inside another
+        of its class — the visible pieces of a half hidden object reported
+        beside the whole of it — into one box. Off by default.
         """
         if source is None:
             raise ValueError("predict() needs a source (image, folder, video, url, camera index)")
@@ -152,6 +156,7 @@ class Detector:
             max_det=max_det,
             classes=classes,
             iou=iou,
+            contain=contain,
             save=save,
             show=show,
             project=project,
@@ -168,6 +173,7 @@ class Detector:
         source: Any = None,
         conf: float = 0.25,
         iou: float | None = 0.7,
+        contain: float | None = None,
         match_iou: float = 0.3,
         max_age: int = 30,
         persist: bool = False,
@@ -176,7 +182,7 @@ class Detector:
     ) -> Any:
         """Predict, and keep an id on each box across frames (``boxes.id``).
 
-        ``iou`` is the duplicate filter, as in :meth:`predict`. ``match_iou``
+        ``iou`` and ``contain`` filter duplicates, as in :meth:`predict`. ``match_iou``
         is how much a box must overlap a track in the previous frame to carry
         its id on; ``max_age`` is how many frames a track waits for its object.
         """
@@ -185,7 +191,7 @@ class Detector:
         if self.tracker is None or not persist:
             self.tracker = IoUTracker(iou=match_iou, max_age=max_age)
         kwargs.setdefault("name", "track")
-        gen = self._run(source, conf=conf, iou=iou, tracker=self.tracker,
+        gen = self._run(source, conf=conf, iou=iou, contain=contain, tracker=self.tracker,
                         **self._track_kwargs(kwargs))
         return _Stream(gen, "track") if stream else list(gen)
 
@@ -209,6 +215,7 @@ class Detector:
         max_det: int = 300,
         classes: list[int] | None = None,
         iou: float | None = 0.7,
+        contain: float | None = None,
         save: bool = False,
         show: bool = False,
         project: str = "runs",
@@ -230,7 +237,7 @@ class Detector:
         try:
             for frame in loader:
                 det, speed = predictor(frame.img, conf=conf, max_det=max_det, classes=classes,
-                                       iou=iou)
+                                       iou=iou, contain=contain)
                 if tracker is not None:
                     det = tracker.update(det)
                 result = Results(
@@ -349,7 +356,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 +369,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.3}/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()]
@@ -84,6 +158,52 @@ def drop_duplicates(xyxy: np.ndarray, iou: float = IOU) -> np.ndarray:
     return np.nonzero(alive)[0]
 
 
+#: How much less sure an enclosing box may be than the box inside it and still
+#: count as the whole object rather than a loose box around several.
+CONTAIN_MARGIN = 0.1
+
+
+def drop_contained(xyxy: np.ndarray, scores: np.ndarray, cls: np.ndarray,
+                   contain: float) -> np.ndarray:
+    """Indices of the boxes to keep when a box lies inside another of its class.
+
+    IoU cannot see a part inside a whole: a chair half hidden behind a person
+    can come back as the whole chair (0.64) plus its two visible pieces (0.67,
+    0.61), and each piece overlaps the whole by an IoU of only its share of the
+    area. Measured as the overlap over the *smaller* box instead, a piece sits
+    near 1.0, so a pair of one class sharing at least ``contain`` of the
+    smaller box is one object, and one box goes:
+
+    * the inner one, when the enclosing box is about as sure (within
+      ``CONTAIN_MARGIN``) — the whole object and a piece of it;
+    * the enclosing one, when it is much less sure than the box inside — a
+      loose box around a group, which must not erase the members.
+
+    NMS's "keep the best" would not do here: the best box is a piece, and once
+    it has removed the whole, the other piece no longer overlaps anything and
+    stays — two half chairs. Same class only: a cup on a table is not part of
+    the table. The cost is a real object inside another of its class, a child
+    held by an adult, so this is off unless asked for.
+    """
+    x1, y1, x2, y2 = xyxy.T
+    area = np.clip(x2 - x1, 0, None) * np.clip(y2 - y1, 0, None)
+    gone = np.zeros(len(xyxy), bool)
+    for i in range(len(xyxy)):
+        for j in range(i + 1, len(xyxy)):
+            if cls[i] != cls[j]:
+                continue
+            iw = max(0.0, min(x2[i], x2[j]) - max(x1[i], x1[j]))
+            ih = max(0.0, min(y2[i], y2[j]) - max(y1[i], y1[j]))
+            if iw * ih < contain * max(min(area[i], area[j]), 1e-9):
+                continue
+            big, small = (i, j) if area[i] >= area[j] else (j, i)
+            if scores[big] >= scores[small] - CONTAIN_MARGIN:
+                gone[small] = True
+            else:
+                gone[big] = True
+    return np.nonzero(~gone)[0]
+
+
 def _split_outputs(arrays: list[np.ndarray]) -> tuple[np.ndarray, np.ndarray]:
     """``(boxes, scores)`` for the one image in a batch, whatever order they come in."""
     boxes = next((a for a in arrays if a.shape[-1] == 4), None)
@@ -161,12 +281,15 @@ class Predictor:
         max_det: int = 300,
         classes: list[int] | None = None,
         iou: float | None = IOU,
+        contain: float | None = None,
     ) -> np.ndarray:
         """-> ``(N, 6)`` array of ``x1 y1 x2 y2 conf cls`` in pixels.
 
         ``iou`` drops a box that covers a higher-scoring one by more than that
         IoU, whatever the two classes (``None`` keeps them all): see
-        :func:`drop_duplicates`.
+        :func:`drop_duplicates`. ``contain`` (off by default) also drops a box
+        lying mostly inside another of its class: see
+        :func:`drop_contained`.
         """
         if _looks_like_logits(scores):
             scores = 1.0 / (1.0 + np.exp(-scores))
@@ -189,6 +312,9 @@ class Predictor:
         if iou is not None and len(xyxy) > 1:
             kept = drop_duplicates(xyxy, iou)
             xyxy, best, cls = xyxy[kept], best[kept], cls[kept]
+        if contain is not None and len(xyxy) > 1:
+            kept = drop_contained(xyxy, best, cls, contain)
+            xyxy, best, cls = xyxy[kept], best[kept], cls[kept]
         return np.concatenate(
             [xyxy, best.astype(np.float32)[:, None], cls.astype(np.float32)[:, None]], axis=1
         )
@@ -200,6 +326,7 @@ class Predictor:
         max_det: int = 300,
         classes: list[int] | None = None,
         iou: float | None = IOU,
+        contain: float | None = None,
     ) -> tuple[np.ndarray, dict[str, float]]:
         """Run one image; returns ``(detections (N,6), speed dict in ms)``."""
         import time
@@ -209,7 +336,8 @@ class Predictor:
         t1 = time.perf_counter()
         boxes, scores = self.infer(tensor)
         t2 = time.perf_counter()
-        det = self.postprocess(boxes, scores, img.shape[:2], conf, max_det, classes, iou)
+        det = self.postprocess(boxes, scores, img.shape[:2], conf, max_det, classes, iou,
+                               contain)
         t3 = time.perf_counter()
         speed = {
             "preprocess": (t1 - t0) * 1e3,

{easydetect-0.2.1 → easydetect-0.2.3}/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.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "easydetect"
-version = "0.2.1"
+version = "0.2.3"
 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.3}/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.3}/tests/test_predictor.py

@@ -147,3 +147,57 @@ def test_one_predictor_can_be_shared_between_threads(tiny_ir):
     for thread in threads:
         thread.join()
     assert failures == []
+
+
+def test_a_half_hidden_object_and_its_pieces_become_one_box():
+    from easydetect.predictor import drop_contained, drop_duplicates
+
+    # best first: a piece (0.67), the whole chair (0.64), another piece (0.61),
+    # and a second chair beside it that only touches the first
+    xyxy = np.array([[320, 520, 372, 640],     # lower right piece
+                     [285, 455, 373, 640],     # the whole chair
+                     [285, 470, 330, 640],     # left piece
+                     [380, 455, 470, 640]],    # another chair
+                    np.float32)
+    scores = np.array([0.67, 0.64, 0.61, 0.60])
+    cls = np.array([56, 56, 56, 56])
+    assert list(drop_duplicates(xyxy, 0.7)) == [0, 1, 2, 3]     # IoU sees four chairs
+    # the whole chair stays, its pieces go; the chair beside it is left alone
+    assert list(drop_contained(xyxy, scores, cls, 0.8)) == [1, 3]
+
+
+def test_a_loose_box_around_a_group_does_not_erase_its_members():
+    from easydetect.predictor import drop_contained
+
+    xyxy = np.array([[0, 0, 50, 100], [50, 0, 100, 100], [0, 0, 100, 100]], np.float32)
+    scores = np.array([0.9, 0.85, 0.3])               # two people, and a box around both
+    assert list(drop_contained(xyxy, scores, np.zeros(3), 0.8)) == [0, 1]
+
+
+def test_containment_never_crosses_classes():
+    from easydetect.predictor import drop_contained
+
+    xyxy = np.array([[0, 0, 100, 100], [10, 10, 40, 40]], np.float32)   # a cup on a table
+    scores = np.array([0.8, 0.7])
+    assert list(drop_contained(xyxy, scores, np.array([60, 41]), 0.8)) == [0, 1]
+    assert list(drop_contained(xyxy, scores, np.array([60, 60]), 0.8)) == [0]
+
+
+def test_the_whole_path_merges_pieces_only_when_asked():
+    """postprocess as predict() runs it: model outputs in, boxes out, at 640 x 720."""
+    from easydetect.predictor import Predictor
+
+    pieces = [(320, 520, 372, 640, 0.67), (285, 455, 373, 640, 0.64), (285, 470, 330, 640, 0.61),
+              (40, 100, 340, 640, 0.97)]                   # three chair boxes and a person
+    boxes = np.array([[(x1 + x2) / 2 / 640, (y1 + y2) / 2 / 720, (x2 - x1) / 640,
+                       (y2 - y1) / 720] for x1, y1, x2, y2, _ in pieces], np.float32)
+    scores = np.zeros((4, 80), np.float32)
+    for i, (*_, s) in enumerate(pieces):
+        scores[i, 0 if i == 3 else 56] = s
+    post = Predictor.postprocess.__get__(object.__new__(Predictor))
+    default = post(boxes, scores, (720, 640), conf=0.5)
+    merged = post(boxes, scores, (720, 640), conf=0.5, contain=0.8)
+    assert len(default) == 4                                # NMS at 0.7 keeps every piece
+    assert sorted(merged[:, 5].tolist()) == [0, 56]         # one person, one chair
+    chair = merged[merged[:, 5] == 56][0]
+    assert chair[:4].round().tolist() == [285, 455, 373, 640]   # the whole of it

{easydetect-0.2.1 → easydetect-0.2.3}/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