martin/object-segmentation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Making it installabel package and switching to segmentation mode

Browse Source
This commit is contained in:
Martin Laasmaa
2025-12-05 15:51:16 +02:00
parent 9011276584
commit 310e0b2285
20 changed files with 667 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
src/database/db_manager.py
View File

@@ -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