fix: commit yolo-seg(#6)

Author: PyTs1n9

backend/alembic/versions/5baeca04aab3_initial_schema.py

@@ -11,6 +11,10 @@
 
 from alembic import op
 
+# Use PostgreSQL's efficient JSONB where available, while retaining support for
+# the zero-config SQLite development database.
+JSON_TYPE = sa.JSON().with_variant(postgresql.JSONB(), "postgresql")
+
 # revision identifiers
 revision: str = '5baeca04aab3'
 down_revision: str | None = None
@@ -24,7 +28,7 @@ def upgrade() -> None:
     sa.Column('id', sa.UUID(), nullable=False),
     sa.Column('image_path', sa.Text(), nullable=False),
     sa.Column('image_name', sa.String(length=512), nullable=False),
-    sa.Column('categories', postgresql.JSONB(astext_type=sa.Text()), nullable=False),
+    sa.Column('categories', JSON_TYPE, nullable=False),
     sa.Column('model_name', sa.String(length=256), nullable=False),
     sa.Column('image_width', sa.Integer(), nullable=False),
     sa.Column('image_height', sa.Integer(), nullable=False),
@@ -43,9 +47,9 @@ def upgrade() -> None:
     sa.Column('epochs', sa.Integer(), nullable=False),
     sa.Column('imgsz', sa.Integer(), nullable=False),
     sa.Column('batch', sa.Integer(), nullable=False),
-    sa.Column('class_map', postgresql.JSONB(astext_type=sa.Text()), nullable=True),
+    sa.Column('class_map', JSON_TYPE, nullable=True),
     sa.Column('status', sa.String(length=32), nullable=False),
-    sa.Column('metrics', postgresql.JSONB(astext_type=sa.Text()), nullable=True),
+    sa.Column('metrics', JSON_TYPE, nullable=True),
     sa.Column('model_path', sa.Text(), nullable=True),
     sa.Column('error_message', sa.Text(), nullable=True),
     sa.Column('created_at', sa.DateTime(timezone=True), nullable=False),

backend/alembic/versions/c3d4e5f6a7b8_add_mask_polygon_to_boxes.py

@@ -11,14 +11,17 @@
 
 from alembic import op
 
+
+JSON_TYPE = sa.JSON().with_variant(postgresql.JSONB(), "postgresql")
+
 revision: str = 'c3d4e5f6a7b8'
 down_revision: str | None = 'b2c3d4e5f6a7'
 branch_labels: str | Sequence[str] | None = None
 depends_on: str | Sequence[str] | None = None
 
 
 def upgrade() -> None:
-    op.add_column('detection_boxes', sa.Column('mask_polygon', postgresql.JSONB(), nullable=True))
+    op.add_column('detection_boxes', sa.Column('mask_polygon', JSON_TYPE, nullable=True))
 
 
 def downgrade() -> None:

backend/app/api/routes/export.py

@@ -9,8 +9,8 @@
 from starlette.responses import Response
 
 from ...core.database import get_db
-from ...schemas.detection import ExportBatchIn
-from ...services.export import FORMAT_LABELS, export_batch, export_single
+from ...schemas.detection import ExportAllIn, ExportBatchIn
+from ...services.export import FORMAT_LABELS, export_all, export_batch, export_single
 
 logger = logging.getLogger(__name__)
 router = APIRouter(prefix="/api/v1", tags=["export"])
@@ -47,3 +47,21 @@ def download_batch_yolo(
         media_type="application/zip",
         headers={"Content-Disposition": f'attachment; filename="{label}_dataset.zip"'},
     )
+
+
+@router.post("/detections/export-all")
+def download_all_detections(
+    body: ExportAllIn,
+    db: Session = Depends(get_db),
+) -> Response:
+    """Download every saved detection as one dataset archive."""
+    try:
+        data = export_all(db, format=body.format)
+    except ValueError as exc:
+        raise HTTPException(400, detail=str(exc)) from exc
+    label = FORMAT_LABELS.get(body.format, body.format)
+    return Response(
+        content=data,
+        media_type="application/zip",
+        headers={"Content-Disposition": f'attachment; filename="{label}_all_dataset.zip"'},
+    )

backend/app/core/gpu_memory.py

@@ -73,7 +73,9 @@ def resolve_max_long_side(self) -> int:
             return 1333
         if total_mb >= 12 * 1024:
             return 1024
-        return 800  # < 12 GB, conservative
+        if total_mb >= 10 * 1024:
+            return 800
+        return 640  # preserve small-object detail on 8 GB cards
 
     def resolve_attn_impl(self) -> str:
         try:

backend/app/repositories/detection.py

@@ -6,12 +6,17 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
+import uuid
 
 from sqlalchemy.orm import Session
 
 from ..models.detection import Detection, DetectionBox
 
 
+def _as_uuid(value: str | uuid.UUID) -> uuid.UUID:
+    return value if isinstance(value, uuid.UUID) else uuid.UUID(value)
+
+
 class DetectionRepository:
     def __init__(self, db: Session) -> None:
         self.db = db
@@ -46,12 +51,13 @@ def create(
         return det
 
     def add_boxes(
-        self, detection_id: str, boxes: list[dict], commit: bool = False
+        self, detection_id: str | uuid.UUID, boxes: list[dict], commit: bool = False
     ) -> list[DetectionBox]:
         """boxes: [{"x1","y1","x2","y2","class_name","confidence","mask_polygon"}, ...]"""
+        detection_uuid = _as_uuid(detection_id)
         entities = [
             DetectionBox(
-                detection_id=detection_id,
+                detection_id=detection_uuid,
                 class_name=b["class_name"],
                 x1=b["x1"],
                 y1=b["y1"],
@@ -70,18 +76,19 @@ def add_boxes(
         return entities
 
     def replace_boxes(
-        self, detection_id: str, boxes: list[dict], commit: bool = False
+        self, detection_id: str | uuid.UUID, boxes: list[dict], commit: bool = False
     ) -> list[DetectionBox]:
         """Delete all existing boxes for a detection and insert new ones."""
+        detection_uuid = _as_uuid(detection_id)
         self.db.query(DetectionBox).filter(
-            DetectionBox.detection_id == detection_id,
+            DetectionBox.detection_id == detection_uuid,
         ).delete()
         # Note: If commit=True, the flush on delete is skipped, and the subsequent
         # add_boxes(commit=True) will commit both the delete and the inserts
         # in a single transaction.
         if not commit:
             self.db.flush()
-        return self.add_boxes(detection_id, boxes, commit=commit)
+        return self.add_boxes(detection_uuid, boxes, commit=commit)
 
     def update_detection(self, detection: Detection, commit: bool = False) -> None:
         if commit:
@@ -99,8 +106,8 @@ def update_box(self, box: DetectionBox, commit: bool = False) -> None:
 
     # ── read ───────────────────────────────────────────
 
-    def get_by_id(self, detection_id: str) -> Detection | None:
-        return self.db.query(Detection).filter(Detection.id == detection_id).first()
+    def get_by_id(self, detection_id: str | uuid.UUID) -> Detection | None:
+        return self.db.query(Detection).filter(Detection.id == _as_uuid(detection_id)).first()
 
     def list(
         self,
@@ -131,8 +138,8 @@ def get_box(self, detection_id: str, box_id: str) -> DetectionBox | None:
         return (
             self.db.query(DetectionBox)
             .filter(
-                DetectionBox.id == box_id,
-                DetectionBox.detection_id == detection_id,
+                DetectionBox.id == _as_uuid(box_id),
+                DetectionBox.detection_id == _as_uuid(detection_id),
             )
             .first()
         )

backend/app/schemas/detection.py

@@ -90,3 +90,7 @@ class DetectionOut(DetectionListItem):
 class ExportBatchIn(BaseSchema):
     detection_ids: list[str] = Field(..., min_length=1)
     format: str = "yolo"
+
+
+class ExportAllIn(BaseSchema):
+    format: str = "yolo-seg"

backend/app/services/export.py

@@ -43,6 +43,16 @@ def export_single_zip(db: Session, detection_id: str, format: str = "yolo") -> b
     return export_batch(db, [detection_id], format=format)
 
 
+def export_all(db: Session, format: str = "yolo-seg") -> bytes:
+    """Export every detection in the database as one dataset archive."""
+    from ..models.detection import Detection
+
+    dets = db.query(Detection).order_by(Detection.created_at, Detection.id).all()
+    if not dets:
+        raise ValueError("No detections available to export")
+    return _export_detections(dets, format)
+
+
 def export_batch(db: Session, detection_ids: list[str], format: str = "yolo") -> bytes:
     """Export multiple detections as a zip file in the requested format."""
     from ..models.detection import Detection
@@ -53,6 +63,10 @@ def export_batch(db: Session, detection_ids: list[str], format: str = "yolo") ->
         if det:
             dets.append(det)
 
+    return _export_detections(dets, format)
+
+
+def _export_detections(dets: list[Detection], format: str) -> bytes:
     unified_map = _build_class_map(dets)
 
     if format == "yolo":

backend/app/services/locate_anything.py

@@ -29,9 +29,15 @@ def _ensure_decord_stub():
     if _decord_loaded:
         return
     if sys.platform in ("darwin", "win32"):
-        stub_dir = str(Path(settings.resolved_model_dir) / "stubs")
-        if Path(stub_dir).exists() and stub_dir not in sys.path:
-            sys.path.insert(0, stub_dir)
+        try:
+            import decord  # type: ignore[import-not-found]  # noqa: F401
+        except ImportError:
+            import decord_stub
+
+            # Transformers checks optional imports before loading remote model
+            # code. Register the bundled image-only compatibility shim under
+            # the package name it expects.
+            sys.modules["decord"] = decord_stub
     _decord_loaded = True
 
 
@@ -80,17 +86,34 @@ def __init__(self, model_path: str, device: str = "mps", progress_cb=None):
         self._set_progress(progress_cb, "downloading", "model", 50)
         attn_impl = gpu_mem.resolve_attn_impl()
         logger.info("Loading model to %s (attn=%s)...", device, attn_impl)
+        load_kwargs = {}
+        if device == "cuda":
+            total_mb = torch.cuda.get_device_properties(0).total_memory // (1024 * 1024)
+            if total_mb < 10 * 1024:
+                # The BF16 model is over 7 GiB. Keeping it entirely on an 8 GiB
+                # display GPU leaves no room for activations and triggers very
+                # slow WDDM shared-memory paging.
+                load_kwargs = {
+                    "device_map": "auto",
+                    "max_memory": {0: "5GiB", "cpu": "12GiB"},
+                }
+                logger.info("Low-VRAM CUDA detected; enabling CPU/GPU weight offload")
         self.model = AutoModel.from_pretrained(
             model_path,
             torch_dtype=self.dtype,
             trust_remote_code=True,
             attn_implementation=attn_impl,
             low_cpu_mem_usage=True,
+            **load_kwargs,
         )
 
         # ── Move to GPU ──
         self._set_progress(progress_cb, "loading", "gpu", 90)
-        self.model = self.model.to(device).eval()
+        if load_kwargs:
+            self.model.eval()
+            logger.info("Model device map: %s", getattr(self.model, "hf_device_map", {}))
+        else:
+            self.model = self.model.to(device).eval()
 
         logger.info("LocateAnything model ready on %s", device)
         self._set_progress(progress_cb, "loaded", "", 100)
@@ -112,7 +135,7 @@ def _predict(
         self,
         image: Image.Image,
         question: str,
-        max_new_tokens: int = 512,
+        max_new_tokens: int = 128,
     ) -> dict:
         messages = [
             {
@@ -393,6 +416,7 @@ def detect(image_path: str | Path, categories: list[str]) -> dict:
         try:
             result = worker.detect(img, categories)
             raw_text = result["answer"]
+            logger.debug("Raw model response: %r", raw_text)
         except Exception as exc:
             raise InferenceError(f"Model inference failed: {exc}") from exc
 

backend/autolabeling.db

@@ -25,7 +25,9 @@
 PIP = VENV / BIN / ("pip.exe" if IS_WIN else "pip")
 PIDFILE = ROOT / ".cli_pids.json"
 
-BACKEND_PORT = int(os.environ.get("BACKEND_PORT", 8000))
+# 8000 and 8500 are reserved by the local Windows networking stack.
+# Keep the environment override for users who need a different port.
+BACKEND_PORT = int(os.environ.get("BACKEND_PORT", 9000))
 FRONTEND_PORT = int(os.environ.get("FRONTEND_PORT", 5173))
 
 # ── Utilities ──────────────────────────────────────────
@@ -207,10 +209,29 @@ def check_sam3_token():
 def download_vlm():
     step("VLM model (LocateAnything-3B, ~6GB)...")
     model_dir = BACKEND / "model"
-    if model_dir.exists() and any(model_dir.iterdir()):
+    required_files = (
+        "config.json",
+        "model.safetensors.index.json",
+        "model-00001-of-00002.safetensors",
+        "model-00002-of-00002.safetensors",
+        "modeling_locateanything.py",
+        "modeling_qwen2.py",
+        "modeling_vit.py",
+        "preprocessor_config.json",
+        "processing_locateanything.py",
+        "processor_config.json",
+        "special_tokens_map.json",
+        "tokenizer_config.json",
+        "vocab.json",
+    )
+    if (
+        model_dir.exists()
+        and not any(model_dir.glob("*.segments.json"))
+        and all((model_dir / name).is_file() for name in required_files)
+    ):
         ok("already cached")
         return
-    print(f"  Downloading to {model_dir}... (~10-30 min)")
+    print(f"  Downloading missing model files to {model_dir}... (~10-30 min)")
     try:
         run(
             [str(PYTHON), "-c",
@@ -311,7 +332,7 @@ def start_frontend():
         warn(f"Port {FRONTEND_PORT} already in use. Try: FRONTEND_PORT={FRONTEND_PORT + 1} python3 cli.py start")
         return None
     proc = subprocess.Popen(
-        ["pnpm", "dev", "--port", str(FRONTEND_PORT)],
+        ["pnpm.cmd", "dev", "--port", str(FRONTEND_PORT)],
         cwd=FRONTEND,
         env={**os.environ, "VITE_BACKEND_PORT": str(BACKEND_PORT)},
         stdout=subprocess.DEVNULL,