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

Fixing annotations in database

Browse Source
This commit is contained in:
Martin Laasmaa
2026-01-21 08:51:39 +02:00
parent d03ffdc4d0
commit 3c8247b3bc
6 changed files with 589 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

140
src/database/db_manager.py
View File

@@ -40,8 +40,15 @@ class DatabaseManager:
conn = self.get_connection()
try:
# Check if annotations table needs migration
# Pre-schema migrations.
# These must run BEFORE executing schema.sql because schema.sql may
# contain CREATE INDEX statements referencing newly added columns.
#
# 1) Check if annotations table needs migration (may drop an old table)
self._migrate_annotations_table(conn)
# 2) Ensure images table has the required columns (e.g. 'source')
self._migrate_images_table(conn)
conn.commit()
# Read schema file and execute
schema_path = Path(__file__).parent / "schema.sql"
@@ -53,6 +60,19 @@ class DatabaseManager:
finally:
conn.close()
def _migrate_images_table(self, conn: sqlite3.Connection) -> None:
"""Migrate images table to include the 'source' column if missing."""
cursor = conn.cursor()
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='images'")
if not cursor.fetchone():
return
cursor.execute("PRAGMA table_info(images)")
columns = {row[1] for row in cursor.fetchall()}
if "source" not in columns:
cursor.execute("ALTER TABLE images ADD COLUMN source TEXT")
def _migrate_annotations_table(self, conn: sqlite3.Connection) -> None:
"""
Migrate annotations table from old schema (class_name) to new schema (class_id).
@@ -233,6 +253,7 @@ class DatabaseManager:
height: int,
captured_at: Optional[datetime] = None,
checksum: Optional[str] = None,
source: Optional[str] = None,
) -> int:
"""
Add a new image to the database.
@@ -253,10 +274,10 @@ class DatabaseManager:
cursor = conn.cursor()
cursor.execute(
"""
INSERT INTO images (relative_path, filename, width, height, captured_at, checksum)
VALUES (?, ?, ?, ?, ?, ?)
INSERT INTO images (relative_path, filename, width, height, captured_at, checksum, source)
VALUES (?, ?, ?, ?, ?, ?, ?)
""",
(relative_path, filename, width, height, captured_at, checksum),
(relative_path, filename, width, height, captured_at, checksum, source),
)
conn.commit()
return cursor.lastrowid
@@ -286,6 +307,18 @@ class DatabaseManager:
return existing["id"]
return self.add_image(relative_path, filename, width, height)
def set_image_source(self, image_id: int, source: Optional[str]) -> bool:
"""Set/update the source marker for an image row."""
conn = self.get_connection()
try:
cursor = conn.cursor()
cursor.execute("UPDATE images SET source = ? WHERE id = ?", (source, int(image_id)))
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
# ==================== Detection Operations ====================
def add_detection(
@@ -658,6 +691,84 @@ class DatabaseManager:
# ==================== Annotation Operations ====================
def get_images_summary(
self,
name_filter: Optional[str] = None,
source_filter: Optional[str] = None,
order_by: str = "filename",
order_dir: str = "ASC",
limit: Optional[int] = None,
offset: int = 0,
) -> List[Dict]:
"""Return all images with annotation counts (including zero).
This is used by the Annotation tab to populate the image list even when
no annotations exist yet.
Args:
name_filter: Optional substring filter applied to filename/relative_path.
order_by: One of: 'filename', 'relative_path', 'annotation_count', 'added_at'.
order_dir: 'ASC' or 'DESC'.
limit: Optional max number of rows.
offset: Pagination offset.
Returns:
List of dicts: {id, relative_path, filename, added_at, annotation_count}
"""
allowed_order_by = {
"filename": "i.filename",
"relative_path": "i.relative_path",
"annotation_count": "annotation_count",
"added_at": "i.added_at",
}
order_expr = allowed_order_by.get(order_by, "i.filename")
dir_norm = str(order_dir).upper().strip()
if dir_norm not in {"ASC", "DESC"}:
dir_norm = "ASC"
conn = self.get_connection()
try:
params: List[Any] = []
where_clauses: List[str] = []
if name_filter:
token = f"%{name_filter}%"
where_clauses.append("(i.filename LIKE ? OR i.relative_path LIKE ?)")
params.extend([token, token])
if source_filter:
where_clauses.append("i.source = ?")
params.append(source_filter)
where_sql = ""
if where_clauses:
where_sql = "WHERE " + " AND ".join(where_clauses)
limit_sql = ""
if limit is not None:
limit_sql = " LIMIT ? OFFSET ?"
params.extend([int(limit), int(offset)])
query = f"""
SELECT
i.id,
i.relative_path,
i.filename,
i.added_at,
COUNT(a.id) AS annotation_count
FROM images i
LEFT JOIN annotations a ON a.image_id = i.id
{where_sql}
GROUP BY i.id
ORDER BY {order_expr} {dir_norm}
{limit_sql}
"""
cursor = conn.cursor()
cursor.execute(query, params)
return [dict(row) for row in cursor.fetchall()]
finally:
conn.close()
def get_annotated_images_summary(
self,
name_filter: Optional[str] = None,
@@ -832,6 +943,27 @@ class DatabaseManager:
finally:
conn.close()
def delete_annotations_for_image(self, image_id: int) -> int:
"""Delete ALL annotations for a specific image.
This is primarily used for import/overwrite workflows.
Args:
image_id: ID of the image whose annotations should be deleted.
Returns:
Number of rows deleted.
"""
conn = self.get_connection()
try:
cursor = conn.cursor()
cursor.execute("DELETE FROM annotations WHERE image_id = ?", (int(image_id),))
conn.commit()
return int(cursor.rowcount or 0)
finally:
conn.close()
# ==================== Object Class Operations ====================
def get_object_classes(self) -> List[Dict]:

8
src/database/schema.sql
View File

@@ -7,7 +7,7 @@ CREATE TABLE IF NOT EXISTS models (
model_name TEXT NOT NULL,
model_version TEXT NOT NULL,
model_path TEXT NOT NULL,
base_model TEXT NOT NULL DEFAULT 'yolov8s.pt',
base_model TEXT NOT NULL DEFAULT 'yolo11s.pt',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
training_params TEXT, -- JSON string of training parameters
metrics TEXT, -- JSON string of validation metrics
@@ -23,7 +23,8 @@ CREATE TABLE IF NOT EXISTS images (
height INTEGER,
captured_at TIMESTAMP,
added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
checksum TEXT
checksum TEXT,
source TEXT
);
-- Detections table: stores detection results
@@ -82,7 +83,8 @@ CREATE INDEX IF NOT EXISTS idx_detections_detected_at ON detections(detected_at)
CREATE INDEX IF NOT EXISTS idx_detections_confidence ON detections(confidence);
CREATE INDEX IF NOT EXISTS idx_images_relative_path ON images(relative_path);
CREATE INDEX IF NOT EXISTS idx_images_added_at ON images(added_at);
CREATE INDEX IF NOT EXISTS idx_images_source ON images(source);
CREATE INDEX IF NOT EXISTS idx_annotations_image_id ON annotations(image_id);
CREATE INDEX IF NOT EXISTS idx_annotations_class_id ON annotations(class_id);
CREATE INDEX IF NOT EXISTS idx_models_created_at ON models(created_at);
CREATE INDEX IF NOT EXISTS idx_object_classes_class_name ON object_classes(class_name);
CREATE INDEX IF NOT EXISTS idx_object_classes_class_name ON object_classes(class_name);