magiclabel 0.1.0__py3-none-any.whl

magiclabel/__init__.py

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

magiclabel/annotator.py

@@ -0,0 +1,85 @@
+import cv2
+from pathlib import Path
+from typing import List, Optional
+
+from rich.progress import Progress
+
+DEFAULT_CLOSED_MODEL = "yolov8n.pt"        # closed-set: 80 COCO classes
+DEFAULT_OPEN_VOCAB_MODEL = "yolov8s-world.pt"  # open-vocabulary: detect anything from text
+
+
+def detect_images(
+    source: Path,
+    output: Path,
+    conf: float,
+    prompt_classes: Optional[List[str]] = None,
+    model_name: Optional[str] = None,
+) -> dict:
+    """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).
+      * open-vocabulary: when `prompt_classes` is given, YOLO-World detects those
+        named objects with no training (e.g. ["missile", "launcher"]).
+
+    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}.
+    """
+    output = Path(output)
+    output.mkdir(parents=True, exist_ok=True)
+    images = _get_image_files(source)
+    if not images:
+        raise FileNotFoundError(f"No images found in {source}")
+
+    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:
+            raise ValueError("No class names given in --prompt")
+        detector.set_classes(class_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)]
+
+    total_boxes = 0
+    with Progress() as progress:
+        task = progress.add_task("Detecting...", total=len(images))
+        for img_path in 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()
+                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))
+                h, w = img.shape[:2]
+                _save_yolo_boxes(output / f"{img_path.stem}.txt", abs_boxes, cls_ids, w, h)
+                total_boxes += len(abs_boxes)
+            progress.advance(task)
+
+    _write_classes_file(output, class_names)
+    return {"classes": class_names, "images": len(images), "boxes": total_boxes}
+
+
+def _save_yolo_boxes(txt_path, boxes, classes, img_w, img_h):
+    with open(txt_path, "w") as f:
+        for box, cls in zip(boxes, classes):
+            x, y, w_box, h_box = box
+            xc = (x + w_box / 2) / img_w
+            yc = (y + h_box / 2) / img_h
+            nw = w_box / img_w
+            nh = h_box / 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/cli.py

@@ -0,0 +1,114 @@
+import typer
+from pathlib import Path
+from typing import List, Optional
+from rich.console import Console
+from magiclabel.annotator import detect_images
+from magiclabel.convertor import yolo_to_coco
+from magiclabel.review_app import launch_review
+from magiclabel.trainer import train_model
+from magiclabel.utils import suppress_ultralytics_output
+
+suppress_ultralytics_output()
+console = Console()
+app = typer.Typer(help="AI‑assisted image annotation CLI")
+
+@app.command()
+def detect(
+    source: Path = typer.Argument(..., help="Directory containing images"),
+    prompt: Optional[str] = typer.Option(
+        None, "--prompt",
+        help="Comma-separated object names for open-vocabulary detection, e.g. "
+             "'missile,launcher,radar dish'. No training needed — detects anything you name.",
+    ),
+    model: Optional[str] = typer.Option(
+        None, help="Model file. Default: yolov8n.pt (closed-set), or yolov8s-world.pt with --prompt.",
+    ),
+    output: Path = typer.Option("./labels_detect", help="Output directory for labels"),
+    conf: float = typer.Option(0.25, help="Confidence threshold"),
+):
+    """Auto‑annotate with bounding boxes (closed-set YOLO, or open-vocabulary with --prompt)."""
+    prompt_classes = [c.strip() for c in prompt.split(",")] if prompt else None
+    try:
+        summary = detect_images(source, output, conf, prompt_classes, model)
+    except Exception as e:
+        console.print(f"[red]Error:[/] {e}")
+        raise typer.Exit(1)
+    mode = "open-vocabulary" if prompt_classes else "closed-set"
+    console.print(
+        f"[green]✓[/] {summary['boxes']} boxes over {summary['images']} images "
+        f"saved to {output} ({mode})"
+    )
+    if prompt_classes:
+        console.print(f"  classes: {', '.join(summary['classes'])}")
+
+@app.command()
+def yolo2coco(
+    images_dir: Path = typer.Argument(..., help="Image directory"),
+    labels_dir: Path = typer.Argument(..., help="YOLO labels directory (detection format)"),
+    output_json: Path = typer.Argument(..., help="Output COCO JSON"),
+    class_names: List[str] = typer.Option(..., "--class", help="Class names in order (0..N-1)"),
+):
+    """Convert YOLO detection labels to COCO JSON."""
+    try:
+        summary = yolo_to_coco(images_dir, labels_dir, output_json, list(class_names))
+    except Exception as e:
+        console.print(f"[red]Error:[/] {e}")
+        raise typer.Exit(1)
+    if summary["auto_named"]:
+        ids = ", ".join(str(i) for i in summary["auto_named"])
+        console.print(
+            f"[yellow]![/] No --class name given for class id(s) {ids}; "
+            f"named them 'class_<id>'. Pass a --class name for every id to fix this."
+        )
+    console.print(
+        f"[green]✓[/] COCO JSON saved to {output_json} "
+        f"({summary['images']} images, {summary['annotations']} annotations, "
+        f"{summary['categories']} categories)"
+    )
+
+@app.command()
+def review(
+    images_dir: Path = typer.Argument(..., help="Image directory"),
+    labels_dir: Path = typer.Argument(..., help="Directory with YOLO label .txt files"),
+    model: str = typer.Option("yolov8n.pt", help="Model whose class names label the boxes"),
+):
+    """Launch a Gradio web interface to review and edit annotations."""
+    try:
+        launch_review(str(images_dir), str(labels_dir), model)
+    except Exception as e:
+        console.print(f"[red]Error:[/] {e}")
+        raise typer.Exit(1)
+
+@app.command()
+def train(
+    images_dir: Path = typer.Argument(..., help="Image directory"),
+    labels_dir: Path = typer.Argument(..., help="Directory with YOLO labels + classes.txt"),
+    name: str = typer.Option("custom.pt", help="Name for your trained model, e.g. rohan.pt"),
+    base: str = typer.Option("yolov8n.pt", help="Base model to fine-tune"),
+    epochs: int = typer.Option(50, help="Training epochs"),
+    imgsz: int = typer.Option(640, help="Training image size"),
+    val_split: float = typer.Option(0.2, help="Fraction of images held out for validation"),
+):
+    """Train a custom YOLO model on your labeled dataset → a named .pt you can reuse."""
+    try:
+        summary = train_model(str(images_dir), str(labels_dir), name, base, epochs, imgsz, val_split)
+    except Exception as e:
+        console.print(f"[red]Error:[/] {e}")
+        raise typer.Exit(1)
+    console.print(
+        f"[green]✓[/] Trained [bold]{summary['model']}[/] on "
+        f"{summary['train']} train / {summary['val']} val images ({summary['epochs']} epochs)"
+    )
+    console.print(f"  classes: {', '.join(summary['classes'])}")
+    console.print(f"  Use it:  magiclabel detect ./new_images --model {summary['model']}")
+
+@app.command()
+def export_classes(model: str = typer.Argument("yolov8n.pt")):
+    """Print class names used by a YOLO model."""
+    from ultralytics import YOLO
+    detector = YOLO(model)
+    for k, v in detector.names.items():
+        console.print(f"{k}: {v}")
+
+if __name__ == "__main__":
+    app()

magiclabel/convertor.py

@@ -0,0 +1,111 @@
+import json
+from pathlib import Path
+
+import cv2
+
+from magiclabel.utils import get_image_files
+
+
+def yolo_to_coco(
+    images_dir: Path,
+    labels_dir: Path,
+    output_json: Path,
+    class_names: list,
+):
+    """Convert a YOLO detection dataset (one .txt per image, normalized
+    `cls xc yc w h` lines) into a single COCO-format JSON file.
+
+    COCO category ids are 1-indexed (category_id = yolo_class_id + 1), matching
+    the COCO convention; image and annotation ids are 1-indexed as well. Class
+    ids that have no provided name are auto-named `class_<id>` so conversion
+    always produces valid output.
+
+    Returns a summary dict: {images, annotations, categories, auto_named}.
+    """
+    images_dir = Path(images_dir)
+    labels_dir = Path(labels_dir)
+
+    images = get_image_files(images_dir)
+    if not images:
+        raise FileNotFoundError(f"No images found in {images_dir}")
+
+    # First pass: discover every class id that actually appears in the labels.
+    present = set()
+    for img_path in images:
+        label_path = labels_dir / f"{img_path.stem}.txt"
+        if not label_path.exists():
+            continue
+        with open(label_path) as f:
+            for line in f:
+                parts = line.strip().split()
+                if len(parts) >= 5:
+                    present.add(int(float(parts[0])))
+
+    def name_for(cls_id):
+        if 0 <= cls_id < len(class_names):
+            return class_names[cls_id]
+        return f"class_{cls_id}"
+
+    # Categories cover both the names the user supplied and any id seen in labels.
+    all_ids = sorted(set(range(len(class_names))) | present)
+    categories = [
+        {"id": i + 1, "name": name_for(i), "supercategory": "none"}
+        for i in all_ids
+    ]
+    auto_named = sorted(i for i in present if i >= len(class_names))
+
+    coco = {"images": [], "annotations": [], "categories": categories}
+    ann_id = 1
+
+    for img_id, img_path in enumerate(images, start=1):
+        img = cv2.imread(str(img_path))
+        if img is None:
+            raise RuntimeError(f"Could not read image: {img_path}")
+        h, w = img.shape[:2]
+
+        coco["images"].append({
+            "id": img_id,
+            "file_name": img_path.name,
+            "width": w,
+            "height": h,
+        })
+
+        label_path = labels_dir / f"{img_path.stem}.txt"
+        if not label_path.exists():
+            continue  # image with no objects is still valid
+
+        with open(label_path) as f:
+            for line in f:
+                parts = line.strip().split()
+                if len(parts) < 5:
+                    continue
+                cls = int(float(parts[0]))
+                xc, yc, bw, bh = map(float, parts[1:5])
+
+                abs_w = bw * w
+                abs_h = bh * h
+                x = (xc - bw / 2) * w
+                y = (yc - bh / 2) * h
+
+                coco["annotations"].append({
+                    "id": ann_id,
+                    "image_id": img_id,
+                    "category_id": cls + 1,
+                    "bbox": [round(x, 2), round(y, 2), round(abs_w, 2), round(abs_h, 2)],
+                    "area": round(abs_w * abs_h, 2),
+                    "iscrowd": 0,
+                    "segmentation": [],
+                })
+                ann_id += 1
+
+    output_json = Path(output_json)
+    output_json.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_json, "w") as f:
+        json.dump(coco, f, indent=2)
+
+    return {
+        "images": len(coco["images"]),
+        "annotations": len(coco["annotations"]),
+        "categories": len(categories),
+        "auto_named": auto_named,
+    }

magiclabel/review_app.py

@@ -0,0 +1,212 @@
+import cv2
+from pathlib import Path
+
+import gradio as gr
+from gradio_image_annotation import image_annotator
+
+IMAGE_EXTS = (".jpg", ".jpeg", ".png")
+CLASSES_FILE = "classes.txt"
+
+
+def _image_size(path):
+    img = cv2.imread(str(path))
+    if img is None:
+        raise RuntimeError(f"Could not read image: {path}")
+    h, w = img.shape[:2]
+    return w, h
+
+
+def _model_class_names(model_name):
+    """Return {class_id: name} from a YOLO model, or {} if it can't be loaded."""
+    if not model_name or str(model_name).lower() in ("none", ""):
+        return {}
+    try:
+        from ultralytics import YOLO
+        return {int(k): str(v) for k, v in YOLO(model_name).names.items()}
+    except Exception:
+        return {}
+
+
+def read_yolo_boxes(txt_path, img_w, img_h, name_of=None):
+    """Read a YOLO .txt into image_annotator box dicts; the box label is the class
+    *name* via name_of(id) (defaults to the numeric id as a string)."""
+    name_of = name_of or (lambda c: str(c))
+    boxes = []
+    if not txt_path.exists():
+        return boxes
+    with open(txt_path) as f:
+        for line in f:
+            parts = line.strip().split()
+            if len(parts) < 5:
+                continue
+            try:
+                cls = int(float(parts[0]))
+                xc, yc, w, h = map(float, parts[1:5])
+            except ValueError:
+                continue
+            xmin = int(round((xc - w / 2) * img_w))
+            ymin = int(round((yc - h / 2) * img_h))
+            xmax = int(round((xc + w / 2) * img_w))
+            ymax = int(round((yc + h / 2) * img_h))
+            boxes.append({
+                "xmin": xmin, "ymin": ymin, "xmax": xmax, "ymax": ymax,
+                "label": name_of(cls),
+            })
+    return boxes
+
+
+def load_annotation(img_path, labels_dir, name_of=None):
+    img_path = Path(img_path)
+    w, h = _image_size(img_path)
+    labels_path = Path(labels_dir) / f"{img_path.stem}.txt"
+    boxes = read_yolo_boxes(labels_path, w, h, name_of)
+    return {"image": str(img_path), "boxes": boxes}
+
+
+def clear_boxes(img_path):
+    """Remove every box from the current image (does not write to disk until Save)."""
+    return ({"image": str(Path(img_path)), "boxes": []},
+            "Cleared all boxes — click Save Changes to write the empty label file.")
+
+
+def _fallback_id(label):
+    try:
+        return int(float(label))
+    except (TypeError, ValueError):
+        return 0
+
+
+def save_annotation(annotation, img_path, labels_dir, id_of=None):
+    """Write boxes back to the YOLO .txt. id_of(label)->int maps a box's (free-text)
+    label to its class id; defaults to parsing the label as a number."""
+    id_of = id_of or _fallback_id
+    if annotation is None:
+        return "Nothing to save."
+    img_path = Path(img_path)
+    w, h = _image_size(img_path)
+    labels_path = Path(labels_dir) / f"{img_path.stem}.txt"
+    boxes = annotation.get("boxes", []) or []
+    written = 0
+    with open(labels_path, "w") as f:
+        for b in boxes:
+            # box corners may come in any order; normalize and skip degenerate boxes
+            x1, x2 = sorted((b["xmin"], b["xmax"]))
+            y1, y2 = sorted((b["ymin"], b["ymax"]))
+            bw, bh = x2 - x1, y2 - y1
+            if bw <= 0 or bh <= 0:
+                continue
+            cls = id_of(b.get("label", ""))
+            xc = (x1 + bw / 2) / w
+            yc = (y1 + bh / 2) / h
+            f.write(f"{cls} {xc:.6f} {yc:.6f} {bw / w:.6f} {bh / h:.6f}\n")
+            written += 1
+    return f"Saved {written} box(es) to {labels_path.name}"
+
+
+def _load_classes_file(labels_dir):
+    f = Path(labels_dir) / CLASSES_FILE
+    if not f.exists():
+        return None
+    names = [ln.strip() for ln in f.read_text().splitlines()]
+    while names and names[-1] == "":
+        names.pop()
+    return names
+
+
+class ClassRegistry:
+    """Maps class names <-> ids and persists the list to labels_dir/classes.txt.
+
+    Seeded from an existing classes.txt if present, else from the model's class names.
+    New names typed in the UI are appended automatically and written back to disk, so
+    classes.txt always reflects the dataset's class list (ready for training).
+    """
+
+    def __init__(self, labels_dir, model_name):
+        self.labels_dir = Path(labels_dir)
+        registry = _load_classes_file(labels_dir)
+        if registry is None:
+            seeded = _model_class_names(model_name)
+            registry = ([seeded.get(i, str(i)) for i in range(max(seeded) + 1)]
+                        if seeded else [])
+        self.names = registry
+        self.name_to_id = {n: i for i, n in enumerate(self.names)}
+        self._persist()
+
+    def name_of(self, cid):
+        return self.names[cid] if 0 <= cid < len(self.names) else str(cid)
+
+    def id_of(self, label):
+        name = str(label).strip() or "object"
+        if name not in self.name_to_id:
+            self.name_to_id[name] = len(self.names)
+            self.names.append(name)
+            self._persist()
+        return self.name_to_id[name]
+
+    def _persist(self):
+        self.labels_dir.mkdir(parents=True, exist_ok=True)
+        (self.labels_dir / CLASSES_FILE).write_text("\n".join(self.names) + "\n")
+
+
+def launch_review(images_dir: str, labels_dir: str, model_name: str = "yolov8n.pt"):
+    images = sorted([str(p) for p in Path(images_dir).glob("*") if p.suffix.lower() in IMAGE_EXTS])
+    if not images:
+        raise RuntimeError("No images found")
+
+    labels_dir = str(labels_dir)
+    registry = ClassRegistry(labels_dir, model_name)
+
+    # Open on the first image that actually has boxes, so the view never looks empty
+    # when labels exist elsewhere in the folder (falls back to the first image).
+    def _has_boxes(img):
+        txt = Path(labels_dir) / f"{Path(img).stem}.txt"
+        return txt.exists() and any(ln.strip() for ln in txt.read_text().splitlines())
+
+    start_image = next((im for im in images if _has_boxes(im)), images[0])
+
+    # labels_dir / registry are constant for the session — bind via closures rather than
+    # gr.State (a State value is not reliably delivered to handlers over the API).
+    def _load(img_path):
+        return load_annotation(img_path, labels_dir, registry.name_of)
+
+    def _clear(img_path):
+        return clear_boxes(img_path)
+
+    def _save(annotation, img_path):
+        msg = save_annotation(annotation, img_path, labels_dir, registry.id_of)
+        return load_annotation(img_path, labels_dir, registry.name_of), msg
+
+    with gr.Blocks(title="MagicLabel Review") as demo:
+        gr.Markdown("## Review & Fix Annotations")
+        gr.Markdown(
+            "**Drag on the image to draw a new box** for objects the model missed. "
+            "Click a box to move or resize it via its handles, or delete it with the box's remove "
+            "button.\n\n"
+            "**To set a box's class:** double‑click the box to open its editor and **type any class "
+            "name you want** (e.g. `missile`, `crack`, `logo`) in the label field, then click the "
+            "apply/✓ button to commit it. New names are added to your class list automatically and "
+            f"saved to `{CLASSES_FILE}` in the labels folder — that file is the class list you'll "
+            "train on. After **Save Changes**, the view reloads from the saved file so you can "
+            "confirm exactly what was written."
+        )
+
+        image_dropdown = gr.Dropdown(choices=images, label="Select Image", value=start_image)
+        annotator = image_annotator(
+            value=_load(start_image),
+            label="Annotations — drag to draw, double‑click a box to type a class name",
+            box_thickness=2,
+            image_type="filepath",
+            show_remove_button=True,
+            enable_keyboard_shortcuts=True,
+            disable_edit_boxes=False,
+        )
+        with gr.Row():
+            clear_btn = gr.Button("Clear All Boxes")
+            save_btn = gr.Button("Save Changes", variant="primary")
+        status = gr.Textbox(label="Status", interactive=False)
+
+        image_dropdown.change(fn=_load, inputs=[image_dropdown], outputs=annotator)
+        clear_btn.click(fn=_clear, inputs=[image_dropdown], outputs=[annotator, status])
+        save_btn.click(fn=_save, inputs=[annotator, image_dropdown], outputs=[annotator, status])
+
+    demo.launch(server_name="127.0.0.1", server_port=7860, share=False)

magiclabel/trainer.py

@@ -0,0 +1,110 @@
+import shutil
+from pathlib import Path
+
+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
+
+
+def _split(labeled, val_split):
+    """Deterministic train/val split: every Nth labeled image goes to validation."""
+    if val_split <= 0:
+        return labeled, labeled[:1]
+    step = max(2, round(1 / val_split))
+    train_set = [p for i, p in enumerate(labeled) if i % step != 0]
+    val_set = [p for i, p in enumerate(labeled) if i % step == 0]
+    if not val_set:
+        val_set.append(train_set.pop())
+    if not train_set:
+        train_set.append(val_set[0])
+    return train_set, val_set
+
+
+def _build_dataset(images_dir, labels_dir, names, train_set, val_set, work_dir):
+    if work_dir.exists():
+        shutil.rmtree(work_dir)
+    for split, items in (("train", train_set), ("val", val_set)):
+        (work_dir / "images" / split).mkdir(parents=True, exist_ok=True)
+        (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")
+
+    # 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")
+    return work_dir / "data.yaml"
+
+
+def train_model(
+    images_dir: str,
+    labels_dir: str,
+    name: str,
+    base: str = "yolov8n.pt",
+    epochs: int = 50,
+    imgsz: int = 640,
+    val_split: float = 0.2,
+) -> 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.
+    """
+    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:
+        raise ValueError(
+            f"Need at least 2 labeled images to train; found {len(labeled)}. "
+            "Label more images in `review` first."
+        )
+
+    train_set, val_set = _split(labeled, val_split)
+
+    out_path = Path(name)
+    if out_path.suffix != ".pt":
+        out_path = out_path.with_suffix(".pt")
+
+    work_dir = Path(f"{out_path.stem}_dataset").resolve()
+    data_yaml = _build_dataset(images_dir, labels_dir, names, train_set, val_set, work_dir)
+
+    from ultralytics import YOLO
+    model = YOLO(base)
+    model.train(
+        data=str(data_yaml),
+        epochs=epochs,
+        imgsz=imgsz,
+        project=str(work_dir / "runs"),
+        name="train",
+        exist_ok=True,
+        verbose=False,
+    )
+
+    best = work_dir / "runs" / "train" / "weights" / "best.pt"
+    if not best.exists():
+        raise RuntimeError("Training finished but best.pt was not produced.")
+    shutil.copy(best, out_path)
+
+    return {
+        "model": str(out_path),
+        "classes": names,
+        "train": len(train_set),
+        "val": len(val_set),
+        "epochs": epochs,
+        "dataset": str(work_dir),
+    }

magiclabel/utils.py

@@ -0,0 +1,12 @@
+import os
+from pathlib import Path
+from typing import List
+
+def get_image_files(source: Path) -> List[Path]:
+    exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}
+    return sorted([p for p in source.iterdir() if p.suffix.lower() in exts])
+
+def suppress_ultralytics_output():
+    """Reduce Ultralytics logging verbosity."""
+    os.environ["YOLO_VERBOSE"] = "False"
+    os.environ["ULTRALYTICS_SILENT"] = "true"

magiclabel-0.1.0.dist-info/METADATA

@@ -0,0 +1,304 @@
+Metadata-Version: 2.4
+Name: magiclabel
+Version: 0.1.0
+Summary: AI-assisted image annotation CLI for computer vision datasets
+Author: Rohan Rustagi
+License: MIT
+Keywords: annotation,labeling,computer-vision,yolo,dataset,object-detection
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: ultralytics>=8.0.0
+Requires-Dist: typer>=0.9.0
+Requires-Dist: rich
+Requires-Dist: opencv-python
+Requires-Dist: gradio
+Requires-Dist: gradio_image_annotation
+Requires-Dist: torch
+Requires-Dist: torchvision
+Dynamic: license-file
+
+# MagicLabel
+
+> AI‑assisted image labeling CLI — go from raw images to your own trained model, end to end.
+
+Good models start with good data, and labeling that data by hand is the slowest part of any computer‑vision project. **MagicLabel** runs the whole loop from one command line:
+
+1. **Auto‑detect** objects in a folder of images — including custom objects you describe in plain words (no training needed).
+2. **Review & fix** the labels in a browser — draw, move, resize, relabel, delete.
+3. **Train** a custom YOLO model on your labeled data.
+4. **Reuse** your named model (`rohan.pt`, `findelon.pt`, …) to auto‑label the next batch — the loop closes.
+
+Everything runs **locally and free** — no account, no cloud upload, no Docker.
+
+![Python](https://img.shields.io/badge/python-3.8%2B-blue)
+![License](https://img.shields.io/badge/license-MIT-green)
+
+---
+
+## Features
+
+- 🔍 **Auto‑detect** — label an image folder with a YOLO model (the 80 COCO classes out of the box).
+- 🗣️ **Open‑vocabulary detect** — `--prompt "missile,launcher"` detects *anything you name*, with no training (YOLO‑World).
+- 🖊️ **Visual review** — a Gradio web app to draw / move / resize / relabel / delete boxes and save back to YOLO format. Type **any class name you want**.
+- 🎓 **Train** — fine‑tune a custom model on your labeled data and save it as a named `.pt`.
+- 🔄 **Export** — convert YOLO labels to COCO JSON.
+- 🏷️ **Inspect** — print the class names a YOLO model knows.
+
+---
+
+## Who is this for?
+
+Anyone who needs a labeled image dataset and a model to train on it:
+
+- **ML engineers & researchers** bootstrapping a dataset from raw images.
+- **Hobbyists & students** learning the full CV pipeline without paying for a labeling service.
+- **Privacy‑sensitive / air‑gapped users** (medical, defense, proprietary data) who can't upload images to a cloud tool.
+
+---
+
+## Requirements
+
+- **Python 3.8+**
+- Disk space for model weights (downloaded automatically on first use)
+- Works on **CPU**; a CUDA or Apple‑Silicon (MPS) GPU is used automatically when available and is much faster for training
+
+> **Heads‑up:** MagicLabel depends on `torch` and `ultralytics`, so installation pulls ~1 GB of packages. That's normal for a computer‑vision tool.
+
+---
+
+## Installation
+
+### From PyPI
+
+```bash
+pip install magiclabel
+```
+
+### From source
+
+```bash
+git clone https://github.com/your-username/magiclabel.git
+cd magiclabel
+python -m venv .venv
+source .venv/bin/activate        # Windows: .venv\Scripts\activate
+pip install -e .
+```
+
+Verify:
+
+```bash
+magiclabel --help
+```
+
+---
+
+## Quick start — the full loop
+
+```bash
+# 1. Auto-label a folder of images (closed-set, 80 COCO classes)
+magiclabel detect ./images --output ./labels
+
+#    …or detect a CUSTOM object by describing it (no training):
+magiclabel detect ./images --prompt "missile, launcher" --output ./labels
+
+# 2. Review and fix the labels in your browser
+magiclabel review ./images ./labels
+
+# 3. Train your own model on the corrected dataset
+magiclabel train ./images ./labels --name rohan.pt --epochs 50
+
+# 4. Reuse your model to auto-label the next batch
+magiclabel detect ./new_images --model rohan.pt --output ./labels2
+```
+
+---
+
+## Commands
+
+Run `magiclabel <command> --help` at any time for the full option list.
+
+### `detect` — auto‑annotate with bounding boxes
+
+Runs a model over every image in a directory and writes one YOLO‑format `.txt` per image, plus a `classes.txt` (the class list).
+
+```bash
+magiclabel detect SOURCE [OPTIONS]
+```
+
+| Option       | Default                                  | Description                                                                 |
+| ------------ | ---------------------------------------- | --------------------------------------------------------------------------- |
+| `SOURCE`     | _(required)_                             | Directory of images (`.jpg`, `.jpeg`, `.png`, `.bmp`, `.tif`)               |
+| `--prompt`   | _(none)_                                 | Comma‑separated object names for **open‑vocabulary** detection, e.g. `"missile,launcher"`. No training needed. |
+| `--model`    | `yolov8n.pt` / `yolov8s-world.pt`        | Model file. Defaults to `yolov8n.pt` (closed‑set), or `yolov8s-world.pt` when `--prompt` is given. Pass your own `.pt` to use a trained model. |
+| `--output`   | `./labels_detect`                        | Output directory for labels + `classes.txt`                                 |
+| `--conf`     | `0.25`                                   | Confidence threshold (0–1); higher = fewer, surer boxes                     |
+
+**Examples**
+
+```bash
+# Closed-set: the 80 COCO classes
+magiclabel detect ./images --conf 0.3 --output ./labels
+
+# Open-vocabulary: detect objects YOLO has never been trained on
+magiclabel detect ./images --prompt "drone, fixed-wing aircraft" --conf 0.1
+
+# Use your own trained model
+magiclabel detect ./images --model rohan.pt --output ./labels
+```
+
+> **Open‑vocabulary is great for common/clear objects and hit‑or‑miss on unusual ones** (it does the easy 70–80%). For hard objects (e.g. missiles), hand‑label a first batch in `review`, `train` a model, then `detect --model your.pt` to auto‑label the rest.
+
+---
+
+### `review` — fix annotations in your browser
+
+Launches a local web app (default: <http://127.0.0.1:7860>) to inspect and correct labels.
+
+```bash
+magiclabel review IMAGES_DIR LABELS_DIR [--model yolov8n.pt]
+```
+
+| Argument / Option | Description                                                                 |
+| ----------------- | --------------------------------------------------------------------------- |
+| `IMAGES_DIR`      | Directory of images                                                         |
+| `LABELS_DIR`      | Directory of YOLO `.txt` labels (and `classes.txt`)                         |
+| `--model`         | Model whose class names label the boxes (default `yolov8n.pt`). Use `--model none` to start a clean custom dataset with no preset names. |
+
+**In the UI you can:**
+
+- **Draw** a new box — click and drag on the image.
+- **Move / resize** — click a box and drag it or its corner handles.
+- **Set the class** — double‑click a box and **type any class name** (`missile`, `crack`, `logo`), then click the editor's apply/✓ to commit. New names are added automatically.
+- **Delete** — the box's remove button, the <kbd>Delete</kbd> key, or **Clear All Boxes**.
+- **Switch images** — the dropdown at the top (it opens on the first image that already has boxes).
+- **Save Changes** — writes boxes back to the `.txt`; the view then reloads from disk so you can confirm exactly what was saved.
+
+**Custom classes & `classes.txt`:** YOLO label files store integer class ids, so MagicLabel keeps a `classes.txt` (one name per line, id = line number) in your labels folder. Whatever name you type is mapped to a stable id behind the scenes, and new names are appended automatically. **That `classes.txt` is the class list you train on.**
+
+> Edits live only in the browser until you click **Save Changes**. Saving an image with no boxes writes an empty label file (a valid "no objects" annotation).
+
+---
+
+### `train` — train your own model
+
+Fine‑tunes a YOLO model on your labeled dataset and saves it as a named `.pt`.
+
+```bash
+magiclabel train IMAGES_DIR LABELS_DIR [OPTIONS]
+```
+
+| Option         | Default       | Description                                            |
+| -------------- | ------------- | ------------------------------------------------------ |
+| `IMAGES_DIR`   | _(required)_  | Directory of images                                    |
+| `LABELS_DIR`   | _(required)_  | Directory of YOLO labels + `classes.txt`               |
+| `--name`       | `custom.pt`   | Output model name, e.g. `rohan.pt`                     |
+| `--base`       | `yolov8n.pt`  | Base model to fine‑tune                                |
+| `--epochs`     | `50`          | Training epochs                                        |
+| `--imgsz`      | `640`         | Training image size                                    |
+| `--val-split`  | `0.2`         | Fraction of images held out for validation             |
+
+**Example**
+
+```bash
+magiclabel train ./images ./labels --name rohan.pt --epochs 100
+# → rohan.pt  (and a rohan_dataset/ folder with the train/val split + data.yaml)
+```
+
+The command builds a proper YOLO dataset (train/val split + `data.yaml`) from your labels, trains, and copies the best weights to your named file. Use it anywhere a model is accepted:
+
+```bash
+magiclabel detect ./new_images --model rohan.pt
+```
+
+> A useful model needs **dozens‑to‑hundreds of labeled images** and enough epochs. On CPU this can be slow — a GPU is much faster. Start small to validate the pipeline, then scale your data.
+
+---
+
+### `yolo2coco` — convert labels to COCO JSON
+
+```bash
+magiclabel yolo2coco IMAGES_DIR LABELS_DIR OUTPUT_JSON --class NAME [--class NAME ...]
+```
+
+| Argument / Option | Description                                                |
+| ----------------- | --------------------------------------------------------- |
+| `IMAGES_DIR`      | Directory of images                                       |
+| `LABELS_DIR`      | Directory of YOLO detection labels                        |
+| `OUTPUT_JSON`     | Path to write the COCO JSON file                          |
+| `--class`         | Class names **in class‑id order** (0, 1, 2, …), repeated  |
+
+```bash
+magiclabel yolo2coco ./images ./labels annotations.json \
+  --class person --class bicycle --class car
+```
+
+---
+
+### `export-classes` — list a model's class names
+
+```bash
+magiclabel export-classes yolov8n.pt
+# 0: person
+# 1: bicycle
+# 2: car
+# ...
+```
+
+---
+
+## Label format
+
+MagicLabel reads and writes the **YOLO detection** format: one `.txt` per image, sharing the image's base name (`dog.jpg` → `dog.txt`). Each line is one box:
+
+```text
+<class_id> <x_center> <y_center> <width> <height>
+```
+
+All four coordinates are **normalized to 0–1** relative to the image size. Example:
+
+```text
+16 0.593846 0.400000 0.250769 0.258462
+```
+
+A `classes.txt` alongside the labels maps each id to a name (id = line number).
+
+---
+
+## Tips & troubleshooting
+
+- **Model downloads:** the first run of a new model downloads its weights into the current directory and reuses them after.
+- **Speed:** detection is fast on CPU; training is much faster on a GPU (PyTorch uses CUDA/MPS automatically).
+- **"No images found":** the source must be a directory containing images with a supported extension (`.jpg`, `.jpeg`, `.png`, `.bmp`, `.tif`, `.tiff`).
+- **Port 7860 already in use:** an old review server is still running. Stop it with:
+  ```bash
+  kill $(lsof -nP -tiTCP:7860 -sTCP:LISTEN)
+  ```
+- **Restart after changes:** the review server does not hot‑reload — stop it (<kbd>Ctrl</kbd>+<kbd>C</kbd>) and run `magiclabel review …` again.
+
+---
+
+## Project structure
+
+```text
+magiclabel/
+├── cli.py          # Typer CLI entry point (detect, review, train, yolo2coco, export-classes)
+├── annotator.py    # YOLO / YOLO-World detection → label files + classes.txt
+├── review_app.py   # Gradio review/annotation web app
+├── trainer.py      # builds the dataset and trains a custom model
+├── convertor.py    # YOLO → COCO conversion
+└── utils.py        # image discovery + helpers
+```
+
+---
+
+## License
+
+Released under the [MIT License](LICENSE).

magiclabel-0.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+magiclabel/__init__.py,sha256=tuRtv5HV8puvzk36ySyUdVbw33UpC9I-KGlFu9Aco_k,101
+magiclabel/annotator.py,sha256=NMKWdhyavtjp8es5QvM-hFPPRc-FVUNB83STkfOyB10,3292
+magiclabel/cli.py,sha256=dDRhGHvFziDBKKUS40qFH2qGRSkTkhHFgWt3IZfUYOU,4947
+magiclabel/convertor.py,sha256=kwRo4c2WC_uAwTnHEEyv1E0lTYcD5neTBpymfnmGnNk,3619
+magiclabel/review_app.py,sha256=3NoW8O03GYhlb1K6TwGvJa9-4-dWyAPzYALEvS6X8fc,8366
+magiclabel/trainer.py,sha256=JJt95f4f2ad0eYcAn52BPWwL8YuMsvmIDCMXfK_xYvk,3841
+magiclabel/utils.py,sha256=8aoyGTQdnG400s45acUvyNKSahGtVCo6YWWQ4ELcApI,418
+magiclabel-0.1.0.dist-info/licenses/LICENSE,sha256=bp4gXAoWuFqpaEwggOuI1vpNMrHzDnTvVgpXSaj0z30,1070
+magiclabel-0.1.0.dist-info/METADATA,sha256=Oq7fb7jR7iwLAkru89saPT9x8YoKLRe-ytc1UtXSLMQ,12431
+magiclabel-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+magiclabel-0.1.0.dist-info/entry_points.txt,sha256=NB73OtQ4Ix7tsCbq6qCJ_SCdfqAMLcPJ6Deq4IFwfXA,50
+magiclabel-0.1.0.dist-info/top_level.txt,sha256=cU2tTDmN9wpssViI1CkPkL30We0hW2lbUDuALAC6RvY,11
+magiclabel-0.1.0.dist-info/RECORD,,

magiclabel-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

magiclabel-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+magiclabel = magiclabel.cli:app

magiclabel-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Rohan Rustagi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

magiclabel-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+magiclabel