Making it installabel package and switching to segmentation mode

src/database/db_manager.py

@@ -6,7 +6,7 @@ Handles all database operations including CRUD operations, queries, and exports.
 import sqlite3
 import json
 from datetime import datetime
-from typing import List, Dict, Optional, Tuple, Any
+from typing import List, Dict, Optional, Tuple, Any, Union
 from pathlib import Path
 import csv
 import hashlib
@@ -56,7 +56,7 @@ class DatabaseManager:
         model_name: str,
         model_version: str,
         model_path: str,
-        base_model: str = "yolov8s.pt",
+        base_model: str = "yolov8s-seg.pt",
         training_params: Optional[Dict] = None,
         metrics: Optional[Dict] = None,
     ) -> int:
@@ -243,6 +243,7 @@ class DatabaseManager:
         class_name: str,
         bbox: Tuple[float, float, float, float],  # (x_min, y_min, x_max, y_max)
         confidence: float,
+        segmentation_mask: Optional[List[List[float]]] = None,
         metadata: Optional[Dict] = None,
     ) -> int:
         """
@@ -254,6 +255,7 @@ class DatabaseManager:
             class_name: Detected object class
             bbox: Bounding box coordinates (normalized 0-1)
             confidence: Detection confidence score
+            segmentation_mask: Polygon coordinates for segmentation [[x1,y1], [x2,y2], ...]
             metadata: Additional metadata
 
         Returns:
@@ -265,8 +267,8 @@ class DatabaseManager:
             x_min, y_min, x_max, y_max = bbox
             cursor.execute(
                 """
-                INSERT INTO detections (image_id, model_id, class_name, x_min, y_min, x_max, y_max, confidence, metadata)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+                INSERT INTO detections (image_id, model_id, class_name, x_min, y_min, x_max, y_max, confidence, segmentation_mask, metadata)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                 """,
                 (
                     image_id,
@@ -277,6 +279,7 @@ class DatabaseManager:
                     x_max,
                     y_max,
                     confidence,
+                    json.dumps(segmentation_mask) if segmentation_mask else None,
                     json.dumps(metadata) if metadata else None,
                 ),
             )
@@ -302,8 +305,8 @@ class DatabaseManager:
                 bbox = det["bbox"]
                 cursor.execute(
                     """
-                    INSERT INTO detections (image_id, model_id, class_name, x_min, y_min, x_max, y_max, confidence, metadata)
-                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    INSERT INTO detections (image_id, model_id, class_name, x_min, y_min, x_max, y_max, confidence, segmentation_mask, metadata)
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                     """,
                     (
                         det["image_id"],
@@ -314,6 +317,11 @@ class DatabaseManager:
                         bbox[2],
                         bbox[3],
                         det["confidence"],
+                        (
+                            json.dumps(det.get("segmentation_mask"))
+                            if det.get("segmentation_mask")
+                            else None
+                        ),
                         (
                             json.dumps(det.get("metadata"))
                             if det.get("metadata")
@@ -385,9 +393,11 @@ class DatabaseManager:
             detections = []
             for row in cursor.fetchall():
                 det = dict(row)
-                # Parse JSON metadata
+                # Parse JSON fields
                 if det.get("metadata"):
                     det["metadata"] = json.loads(det["metadata"])
+                if det.get("segmentation_mask"):
+                    det["segmentation_mask"] = json.loads(det["segmentation_mask"])
                 detections.append(det)
 
             return detections
@@ -538,6 +548,7 @@ class DatabaseManager:
                     "x_max",
                     "y_max",
                     "confidence",
+                    "segmentation_mask",
                     "detected_at",
                 ]
                 writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
@@ -545,6 +556,11 @@ class DatabaseManager:
 
                 for det in detections:
                     row = {k: det[k] for k in fieldnames if k in det}
+                    # Convert segmentation mask list to JSON string for CSV
+                    if row.get("segmentation_mask") and isinstance(
+                        row["segmentation_mask"], list
+                    ):
+                        row["segmentation_mask"] = json.dumps(row["segmentation_mask"])
                     writer.writerow(row)
 
             return True
@@ -580,6 +596,7 @@ class DatabaseManager:
         class_name: str,
         bbox: Tuple[float, float, float, float],
         annotator: str,
+        segmentation_mask: Optional[List[List[float]]] = None,
         verified: bool = False,
     ) -> int:
         """Add manual annotation."""
@@ -589,10 +606,20 @@ class DatabaseManager:
             x_min, y_min, x_max, y_max = bbox
             cursor.execute(
                 """
-                INSERT INTO annotations (image_id, class_name, x_min, y_min, x_max, y_max, annotator, verified)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                INSERT INTO annotations (image_id, class_name, x_min, y_min, x_max, y_max, segmentation_mask, annotator, verified)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                 """,
-                (image_id, class_name, x_min, y_min, x_max, y_max, annotator, verified),
+                (
+                    image_id,
+                    class_name,
+                    x_min,
+                    y_min,
+                    x_max,
+                    y_max,
+                    json.dumps(segmentation_mask) if segmentation_mask else None,
+                    annotator,
+                    verified,
+                ),
             )
             conn.commit()
             return cursor.lastrowid

src/database/models.py

@@ -5,7 +5,7 @@ These dataclasses represent the database entities.
 
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Optional, Dict, Tuple
+from typing import Optional, Dict, Tuple, List
 
 
 @dataclass
@@ -46,6 +46,9 @@ class Detection:
     class_name: str
     bbox: Tuple[float, float, float, float]  # (x_min, y_min, x_max, y_max)
     confidence: float
+    segmentation_mask: Optional[
+        List[List[float]]
+    ]  # List of polygon coordinates [[x1,y1], [x2,y2], ...]
     detected_at: datetime
     metadata: Optional[Dict]
 
@@ -58,6 +61,9 @@ class Annotation:
     image_id: int
     class_name: str
     bbox: Tuple[float, float, float, float]  # (x_min, y_min, x_max, y_max)
+    segmentation_mask: Optional[
+        List[List[float]]
+    ]  # List of polygon coordinates [[x1,y1], [x2,y2], ...]
     annotator: str
     created_at: datetime
     verified: bool

src/database/schema.sql

@@ -37,6 +37,7 @@ CREATE TABLE IF NOT EXISTS detections (
     x_max REAL NOT NULL CHECK(x_max >= 0 AND x_max <= 1),
     y_max REAL NOT NULL CHECK(y_max >= 0 AND y_max <= 1),
     confidence REAL NOT NULL CHECK(confidence >= 0 AND confidence <= 1),
+    segmentation_mask TEXT,  -- JSON string of polygon coordinates [[x1,y1], [x2,y2], ...]
     detected_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
     metadata TEXT,  -- JSON string for additional metadata
     FOREIGN KEY (image_id) REFERENCES images (id) ON DELETE CASCADE,
@@ -52,6 +53,7 @@ CREATE TABLE IF NOT EXISTS annotations (
     y_min REAL NOT NULL CHECK(y_min >= 0 AND y_min <= 1),
     x_max REAL NOT NULL CHECK(x_max >= 0 AND x_max <= 1),
     y_max REAL NOT NULL CHECK(y_max >= 0 AND y_max <= 1),
+    segmentation_mask TEXT,  -- JSON string of polygon coordinates [[x1,y1], [x2,y2], ...]
     annotator TEXT,
     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
     verified BOOLEAN DEFAULT 0,