Fixing annotations in database

src/database/db_manager.py

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