Adding pen tool for annotation
This commit is contained in:
@@ -30,18 +30,48 @@ class DatabaseManager:
|
||||
# Create directory if it doesn't exist
|
||||
Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Read schema file and execute
|
||||
schema_path = Path(__file__).parent / "schema.sql"
|
||||
with open(schema_path, "r") as f:
|
||||
schema_sql = f.read()
|
||||
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
# Check if annotations table needs migration
|
||||
self._migrate_annotations_table(conn)
|
||||
|
||||
# Read schema file and execute
|
||||
schema_path = Path(__file__).parent / "schema.sql"
|
||||
with open(schema_path, "r") as f:
|
||||
schema_sql = f.read()
|
||||
|
||||
conn.executescript(schema_sql)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _migrate_annotations_table(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
Migrate annotations table from old schema (class_name) to new schema (class_id).
|
||||
"""
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Check if annotations table exists
|
||||
cursor.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='annotations'"
|
||||
)
|
||||
if not cursor.fetchone():
|
||||
# Table doesn't exist yet, no migration needed
|
||||
return
|
||||
|
||||
# Check if table has old schema (class_name column)
|
||||
cursor.execute("PRAGMA table_info(annotations)")
|
||||
columns = {row[1]: row for row in cursor.fetchall()}
|
||||
|
||||
if "class_name" in columns and "class_id" not in columns:
|
||||
# Old schema detected, need to migrate
|
||||
print("Migrating annotations table to new schema with class_id...")
|
||||
|
||||
# Drop old annotations table (assuming no critical data since this is a new feature)
|
||||
cursor.execute("DROP TABLE IF EXISTS annotations")
|
||||
conn.commit()
|
||||
print("Old annotations table dropped, will be recreated with new schema")
|
||||
|
||||
def get_connection(self) -> sqlite3.Connection:
|
||||
"""Get database connection with proper settings."""
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
@@ -593,25 +623,38 @@ class DatabaseManager:
|
||||
def add_annotation(
|
||||
self,
|
||||
image_id: int,
|
||||
class_name: str,
|
||||
class_id: int,
|
||||
bbox: Tuple[float, float, float, float],
|
||||
annotator: str,
|
||||
segmentation_mask: Optional[List[List[float]]] = None,
|
||||
verified: bool = False,
|
||||
) -> int:
|
||||
"""Add manual annotation."""
|
||||
"""
|
||||
Add manual annotation.
|
||||
|
||||
Args:
|
||||
image_id: ID of the image
|
||||
class_id: ID of the object class (foreign key to object_classes)
|
||||
bbox: Bounding box coordinates (normalized 0-1)
|
||||
annotator: Name of person/tool creating annotation
|
||||
segmentation_mask: Polygon coordinates for segmentation
|
||||
verified: Whether annotation has been verified
|
||||
|
||||
Returns:
|
||||
ID of the inserted annotation
|
||||
"""
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
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, segmentation_mask, annotator, verified)
|
||||
INSERT INTO annotations (image_id, class_id, x_min, y_min, x_max, y_max, segmentation_mask, annotator, verified)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
image_id,
|
||||
class_name,
|
||||
class_id,
|
||||
x_min,
|
||||
y_min,
|
||||
x_max,
|
||||
@@ -627,15 +670,178 @@ class DatabaseManager:
|
||||
conn.close()
|
||||
|
||||
def get_annotations_for_image(self, image_id: int) -> List[Dict]:
|
||||
"""Get all annotations for an image."""
|
||||
"""
|
||||
Get all annotations for an image with class information.
|
||||
|
||||
Args:
|
||||
image_id: ID of the image
|
||||
|
||||
Returns:
|
||||
List of annotation dictionaries with joined class information
|
||||
"""
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT * FROM annotations WHERE image_id = ?", (image_id,))
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT
|
||||
a.*,
|
||||
c.class_name,
|
||||
c.color as class_color,
|
||||
c.description as class_description
|
||||
FROM annotations a
|
||||
JOIN object_classes c ON a.class_id = c.id
|
||||
WHERE a.image_id = ?
|
||||
ORDER BY a.created_at DESC
|
||||
""",
|
||||
(image_id,),
|
||||
)
|
||||
annotations = []
|
||||
for row in cursor.fetchall():
|
||||
ann = dict(row)
|
||||
if ann.get("segmentation_mask"):
|
||||
ann["segmentation_mask"] = json.loads(ann["segmentation_mask"])
|
||||
annotations.append(ann)
|
||||
return annotations
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
# ==================== Object Class Operations ====================
|
||||
|
||||
def get_object_classes(self) -> List[Dict]:
|
||||
"""
|
||||
Get all object classes.
|
||||
|
||||
Returns:
|
||||
List of object class dictionaries
|
||||
"""
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT * FROM object_classes ORDER BY class_name")
|
||||
return [dict(row) for row in cursor.fetchall()]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def get_object_class_by_id(self, class_id: int) -> Optional[Dict]:
|
||||
"""Get object class by ID."""
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT * FROM object_classes WHERE id = ?", (class_id,))
|
||||
row = cursor.fetchone()
|
||||
return dict(row) if row else None
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def get_object_class_by_name(self, class_name: str) -> Optional[Dict]:
|
||||
"""Get object class by name."""
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT * FROM object_classes WHERE class_name = ?", (class_name,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return dict(row) if row else None
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def add_object_class(
|
||||
self, class_name: str, color: str, description: Optional[str] = None
|
||||
) -> int:
|
||||
"""
|
||||
Add a new object class.
|
||||
|
||||
Args:
|
||||
class_name: Name of the object class
|
||||
color: Hex color code (e.g., '#FF0000')
|
||||
description: Optional description
|
||||
|
||||
Returns:
|
||||
ID of the inserted object class
|
||||
"""
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO object_classes (class_name, color, description)
|
||||
VALUES (?, ?, ?)
|
||||
""",
|
||||
(class_name, color, description),
|
||||
)
|
||||
conn.commit()
|
||||
return cursor.lastrowid
|
||||
except sqlite3.IntegrityError:
|
||||
# Class already exists
|
||||
existing = self.get_object_class_by_name(class_name)
|
||||
return existing["id"] if existing else None
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def update_object_class(
|
||||
self,
|
||||
class_id: int,
|
||||
class_name: Optional[str] = None,
|
||||
color: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Update an object class.
|
||||
|
||||
Args:
|
||||
class_id: ID of the class to update
|
||||
class_name: New class name (optional)
|
||||
color: New color (optional)
|
||||
description: New description (optional)
|
||||
|
||||
Returns:
|
||||
True if updated, False otherwise
|
||||
"""
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
updates = {}
|
||||
if class_name is not None:
|
||||
updates["class_name"] = class_name
|
||||
if color is not None:
|
||||
updates["color"] = color
|
||||
if description is not None:
|
||||
updates["description"] = description
|
||||
|
||||
if not updates:
|
||||
return False
|
||||
|
||||
set_clauses = [f"{key} = ?" for key in updates.keys()]
|
||||
params = list(updates.values()) + [class_id]
|
||||
|
||||
query = f"UPDATE object_classes SET {', '.join(set_clauses)} WHERE id = ?"
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(query, params)
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def delete_object_class(self, class_id: int) -> bool:
|
||||
"""
|
||||
Delete an object class.
|
||||
|
||||
Args:
|
||||
class_id: ID of the class to delete
|
||||
|
||||
Returns:
|
||||
True if deleted, False otherwise
|
||||
"""
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("DELETE FROM object_classes WHERE id = ?", (class_id,))
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
@staticmethod
|
||||
def calculate_checksum(file_path: str) -> str:
|
||||
"""Calculate MD5 checksum of a file."""
|
||||
|
||||
@@ -44,11 +44,27 @@ CREATE TABLE IF NOT EXISTS detections (
|
||||
FOREIGN KEY (model_id) REFERENCES models (id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
-- Annotations table: stores manual annotations (future feature)
|
||||
-- Object classes table: stores annotation class definitions with colors
|
||||
CREATE TABLE IF NOT EXISTS object_classes (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
class_name TEXT NOT NULL UNIQUE,
|
||||
color TEXT NOT NULL, -- Hex color code (e.g., '#FF0000')
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
description TEXT
|
||||
);
|
||||
|
||||
-- Insert default object classes
|
||||
INSERT OR IGNORE INTO object_classes (class_name, color, description) VALUES
|
||||
('cell', '#FF0000', 'Cell object'),
|
||||
('nucleus', '#00FF00', 'Cell nucleus'),
|
||||
('mitochondria', '#0000FF', 'Mitochondria'),
|
||||
('vesicle', '#FFFF00', 'Vesicle');
|
||||
|
||||
-- Annotations table: stores manual annotations
|
||||
CREATE TABLE IF NOT EXISTS annotations (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
image_id INTEGER NOT NULL,
|
||||
class_name TEXT NOT NULL,
|
||||
class_id INTEGER NOT NULL,
|
||||
x_min REAL NOT NULL CHECK(x_min >= 0 AND x_min <= 1),
|
||||
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),
|
||||
@@ -57,7 +73,8 @@ CREATE TABLE IF NOT EXISTS annotations (
|
||||
annotator TEXT,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
verified BOOLEAN DEFAULT 0,
|
||||
FOREIGN KEY (image_id) REFERENCES images (id) ON DELETE CASCADE
|
||||
FOREIGN KEY (image_id) REFERENCES images (id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (class_id) REFERENCES object_classes (id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
-- Create indexes for performance optimization
|
||||
@@ -69,4 +86,6 @@ 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_annotations_image_id ON annotations(image_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_models_created_at ON models(created_at);
|
||||
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);
|
||||
Reference in New Issue
Block a user