magiclabel 0.2.0__tar.gz → 0.2.1__tar.gz

{magiclabel-0.2.0 → magiclabel-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: magiclabel
-Version: 0.2.0
+Version: 0.2.1
 Summary: AI-assisted image annotation CLI for computer vision datasets
 Author: Rohan Rustagi
 License: MIT

{magiclabel-0.2.0 → magiclabel-0.2.1}/magiclabel/__init__.py

@@ -1,2 +1,2 @@
 """magiclabel – Open‑source CLI tool for AI‑assisted image annotation."""
-__version__ = "0.2.0"
+__version__ = "0.2.1"

{magiclabel-0.2.0 → magiclabel-0.2.1}/magiclabel/annotator.py

@@ -4,10 +4,17 @@ from typing import List, Optional
 
 from rich.progress import Progress
 
-DEFAULT_CLOSED_MODEL = "yolov8n.pt"        # closed-set: 80 COCO classes
+DEFAULT_CLOSED_MODEL = "yolov8n.pt"            # closed-set: 80 COCO classes
 DEFAULT_OPEN_VOCAB_MODEL = "yolov8s-world.pt"  # open-vocabulary: detect anything from text
 
 
+def _load_classes(output: Path):
+    f = output / "classes.txt"
+    if not f.exists():
+        return []
+    return [ln.strip() for ln in f.read_text().splitlines() if ln.strip()]
+
+
 def detect_images(
     source: Path,
     output: Path,
@@ -18,13 +25,14 @@ def detect_images(
     """Auto-annotate every image in `source` with bounding boxes.
 
     Two modes:
-      * closed-set (default): a normal YOLO model (e.g. yolov8n.pt) detects its
-        built-in classes (the 80 COCO objects).
+      * closed-set (default): a normal YOLO model detects its built-in classes.
       * open-vocabulary: when `prompt_classes` is given, YOLO-World detects those
-        named objects with no training (e.g. ["missile", "launcher"]).
+        named objects with no training.
 
-    Writes one YOLO `.txt` per image plus a `classes.txt` (the class list, in id
-    order) into `output`. Returns a summary dict {classes, images, boxes}.
+    Class handling (consistent + compact): labels are written by *class name*. The
+    class list lives in `output/classes.txt` and is MERGED, never clobbered — only
+    classes that are actually detected get an id, and re-running detect (or adding
+    custom classes in `review`) keeps ids and names in sync. Returns a summary dict.
     """
     output = Path(output)
     output.mkdir(parents=True, exist_ok=True)
@@ -35,14 +43,25 @@ def detect_images(
     if prompt_classes:
         from ultralytics import YOLOWorld
         detector = YOLOWorld(model_name or DEFAULT_OPEN_VOCAB_MODEL)
-        class_names = [c.strip() for c in prompt_classes if c.strip()]
-        if not class_names:
+        model_names = [c.strip() for c in prompt_classes if c.strip()]
+        if not model_names:
             raise ValueError("No class names given in --prompt")
-        detector.set_classes(class_names)
+        detector.set_classes(model_names)
     else:
         from ultralytics import YOLO
         detector = YOLO(model_name or DEFAULT_CLOSED_MODEL)
-        class_names = [detector.names[i] for i in sorted(detector.names)]
+        model_names = [detector.names[i] for i in sorted(detector.names)]
+
+    # name-consistent registry, seeded from any existing classes.txt (merge, don't clobber)
+    registry = _load_classes(output)
+    name_to_id = {n: i for i, n in enumerate(registry)}
+
+    def reg_id(model_idx):
+        name = model_names[model_idx] if 0 <= model_idx < len(model_names) else str(model_idx)
+        if name not in name_to_id:
+            name_to_id[name] = len(registry)
+            registry.append(name)
+        return name_to_id[name]
 
     total_boxes = 0
     with Progress() as progress:
@@ -51,7 +70,7 @@ def detect_images(
             results = detector.predict(img_path, conf=conf, verbose=False)
             boxes = results[0].boxes
             if boxes is not None and len(boxes) > 0:
-                cls_ids = boxes.cls.int().tolist()
+                cls_ids = [reg_id(mi) for mi in boxes.cls.int().tolist()]
                 xyxy = boxes.xyxy.tolist()
                 abs_boxes = [[x1, y1, x2 - x1, y2 - y1] for x1, y1, x2, y2 in xyxy]
                 img = cv2.imread(str(img_path))
@@ -60,8 +79,8 @@ def detect_images(
                 total_boxes += len(abs_boxes)
             progress.advance(task)
 
-    _write_classes_file(output, class_names)
-    return {"classes": class_names, "images": len(images), "boxes": total_boxes}
+    (output / "classes.txt").write_text("\n".join(registry) + "\n")
+    return {"classes": registry, "images": len(images), "boxes": total_boxes}
 
 
 def _save_yolo_boxes(txt_path, boxes, classes, img_w, img_h):
@@ -75,11 +94,6 @@ def _save_yolo_boxes(txt_path, boxes, classes, img_w, img_h):
             f.write(f"{cls} {xc:.6f} {yc:.6f} {nw:.6f} {nh:.6f}\n")
 
 
-def _write_classes_file(output, class_names):
-    """Write the class list (one name per line, id == line index) for review/training."""
-    (Path(output) / "classes.txt").write_text("\n".join(class_names) + "\n")
-
-
 def _get_image_files(source):
     from magiclabel.utils import get_image_files
     return get_image_files(source)

{magiclabel-0.2.0 → magiclabel-0.2.1}/magiclabel/trainer.py

@@ -5,16 +5,24 @@ from magiclabel.utils import get_image_files
 
 
 def _read_classes(labels_dir: Path):
-    classes_file = labels_dir / "classes.txt"
-    if not classes_file.exists():
-        raise FileNotFoundError(
-            f"{classes_file} not found. Run `magiclabel detect` or `review` first — "
-            "they create the class list (classes.txt) this command trains on."
-        )
-    names = [ln.strip() for ln in classes_file.read_text().splitlines() if ln.strip()]
-    if not names:
-        raise ValueError("classes.txt is empty — there are no classes to train on.")
-    return names
+    f = labels_dir / "classes.txt"
+    if not f.exists():
+        return []
+    return [ln.strip() for ln in f.read_text().splitlines() if ln.strip()]
+
+
+def _used_class_ids(labels_dir: Path, labeled):
+    """All class ids that actually appear across the label files."""
+    used = set()
+    for img in labeled:
+        for ln in (labels_dir / f"{img.stem}.txt").read_text().splitlines():
+            parts = ln.split()
+            if len(parts) >= 5:
+                try:
+                    used.add(int(float(parts[0])))
+                except ValueError:
+                    pass
+    return sorted(used)
 
 
 def _split(labeled, val_split):
@@ -31,7 +39,24 @@ def _split(labeled, val_split):
     return train_set, val_set
 
 
-def _build_dataset(images_dir, labels_dir, names, train_set, val_set, work_dir):
+def _write_remapped_label(src_txt, dst_txt, remap):
+    """Copy a label file, rewriting each box's class id through `remap`."""
+    out = []
+    for ln in src_txt.read_text().splitlines():
+        parts = ln.split()
+        if len(parts) < 5:
+            continue
+        try:
+            old = int(float(parts[0]))
+        except ValueError:
+            continue
+        if old not in remap:
+            continue
+        out.append(" ".join([str(remap[old])] + parts[1:]))
+    dst_txt.write_text("\n".join(out) + ("\n" if out else ""))
+
+
+def _build_dataset(images_dir, labels_dir, names, remap, train_set, val_set, work_dir):
     if work_dir.exists():
         shutil.rmtree(work_dir)
     for split, items in (("train", train_set), ("val", val_set)):
@@ -39,10 +64,12 @@ def _build_dataset(images_dir, labels_dir, names, train_set, val_set, work_dir):
         (work_dir / "labels" / split).mkdir(parents=True, exist_ok=True)
         for img in items:
             shutil.copy(img, work_dir / "images" / split / img.name)
-            shutil.copy(labels_dir / f"{img.stem}.txt",
-                        work_dir / "labels" / split / f"{img.stem}.txt")
+            _write_remapped_label(
+                labels_dir / f"{img.stem}.txt",
+                work_dir / "labels" / split / f"{img.stem}.txt",
+                remap,
+            )
 
-    # data.yaml — quote names so spaces/punctuation are safe
     lines = [f"path: {work_dir.resolve()}", "train: images/train", "val: images/val", "names:"]
     lines += [f'  {i}: "{n}"' for i, n in enumerate(names)]
     (work_dir / "data.yaml").write_text("\n".join(lines) + "\n")
@@ -60,12 +87,13 @@ def train_model(
 ) -> dict:
     """Train a custom YOLO detector on a labeled dataset and save it as a named .pt.
 
-    Reads classes.txt + YOLO label files, builds a train/val split, fine-tunes `base`,
-    and copies the best weights to `name` (e.g. rohan.pt). Returns a summary dict.
+    The dataset is COMPACTED to only the classes that actually appear in the labels:
+    used class ids are remapped to a contiguous 0..K-1, named from classes.txt (with a
+    `class_<id>` fallback for any id missing a name). This keeps training robust even if
+    labels contain gaps or ids beyond classes.txt — no class-count crashes, no bloat.
     """
     images_dir = Path(images_dir)
     labels_dir = Path(labels_dir)
-    names = _read_classes(labels_dir)
 
     labeled = [p for p in get_image_files(images_dir) if (labels_dir / f"{p.stem}.txt").exists()]
     if len(labeled) < 2:
@@ -74,14 +102,21 @@ def train_model(
             "Label more images in `review` first."
         )
 
-    train_set, val_set = _split(labeled, val_split)
+    used = _used_class_ids(labels_dir, labeled)
+    if not used:
+        raise ValueError("No labeled boxes found — every label file is empty.")
+
+    all_names = _read_classes(labels_dir)
+    remap = {old: new for new, old in enumerate(used)}
+    names = [all_names[old] if 0 <= old < len(all_names) else f"class_{old}" for old in used]
 
     out_path = Path(name)
     if out_path.suffix != ".pt":
         out_path = out_path.with_suffix(".pt")
 
+    train_set, val_set = _split(labeled, val_split)
     work_dir = Path(f"{out_path.stem}_dataset").resolve()
-    data_yaml = _build_dataset(images_dir, labels_dir, names, train_set, val_set, work_dir)
+    data_yaml = _build_dataset(images_dir, labels_dir, names, remap, train_set, val_set, work_dir)
 
     from ultralytics import YOLO
     model = YOLO(base)

{magiclabel-0.2.0 → magiclabel-0.2.1}/magiclabel.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: magiclabel
-Version: 0.2.0
+Version: 0.2.1
 Summary: AI-assisted image annotation CLI for computer vision datasets
 Author: Rohan Rustagi
 License: MIT

{magiclabel-0.2.0 → magiclabel-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "magiclabel"
-version = "0.2.0"
+version = "0.2.1"
 description = "AI-assisted image annotation CLI for computer vision datasets"
 readme = "README.md"
 license = {text = "MIT"}