diff --git a/src/database/db_manager.py b/src/database/db_manager.py index 04681a7..5e4dcc4 100644 --- a/src/database/db_manager.py +++ b/src/database/db_manager.py @@ -201,6 +201,28 @@ class DatabaseManager: finally: conn.close() + def delete_model(self, model_id: int) -> bool: + """Delete a model from the database. + + Note: detections referencing this model are deleted automatically via + the `detections.model_id` foreign key (ON DELETE CASCADE). + + Args: + model_id: ID of the model to delete. + + Returns: + True if a model row was deleted, False otherwise. + """ + + conn = self.get_connection() + try: + cursor = conn.cursor() + cursor.execute("DELETE FROM models WHERE id = ?", (model_id,)) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + # ==================== Image Operations ==================== def add_image( diff --git a/src/gui/main_window.py b/src/gui/main_window.py index 3a4df57..811310f 100644 --- a/src/gui/main_window.py +++ b/src/gui/main_window.py @@ -1,6 +1,7 @@ -""" -Main window for the microscopy object detection application. -""" +"""Main window for the microscopy object detection application.""" + +import shutil +from pathlib import Path from PySide6.QtWidgets import ( QMainWindow, @@ -20,6 +21,7 @@ from src.database.db_manager import DatabaseManager from src.utils.config_manager import ConfigManager from src.utils.logger import get_logger from src.gui.dialogs.config_dialog import ConfigDialog +from src.gui.dialogs.delete_model_dialog import DeleteModelDialog from src.gui.tabs.detection_tab import DetectionTab from src.gui.tabs.training_tab import TrainingTab from src.gui.tabs.validation_tab import ValidationTab @@ -91,6 +93,12 @@ class MainWindow(QMainWindow): db_stats_action.triggered.connect(self._show_database_stats) tools_menu.addAction(db_stats_action) + tools_menu.addSeparator() + + delete_model_action = QAction("Delete &Model…", self) + delete_model_action.triggered.connect(self._show_delete_model_dialog) + tools_menu.addAction(delete_model_action) + # Help menu help_menu = menubar.addMenu("&Help") @@ -152,9 +160,7 @@ class MainWindow(QMainWindow): """Center window on screen.""" screen = self.screen().geometry() size = self.geometry() - self.move( - (screen.width() - size.width()) // 2, (screen.height() - size.height()) // 2 - ) + self.move((screen.width() - size.width()) // 2, (screen.height() - size.height()) // 2) def _restore_window_state(self): """Restore window geometry from settings or center window.""" @@ -231,9 +237,229 @@ class MainWindow(QMainWindow): except Exception as e: logger.error(f"Error getting database stats: {e}") - QMessageBox.warning( - self, "Error", f"Failed to get database statistics:\n{str(e)}" - ) + QMessageBox.warning(self, "Error", f"Failed to get database statistics:\n{str(e)}") + + def _show_delete_model_dialog(self) -> None: + """Open the model deletion dialog.""" + dialog = DeleteModelDialog(self.db_manager, self) + if not dialog.exec(): + return + + model_ids = dialog.selected_model_ids + if not model_ids: + return + + self._delete_models(model_ids) + + def _delete_models(self, model_ids: list[int]) -> None: + """Delete one or more models from the database and remove artifacts from disk.""" + + deleted_count = 0 + removed_paths: list[str] = [] + remove_errors: list[str] = [] + + for model_id in model_ids: + model = None + try: + model = self.db_manager.get_model_by_id(int(model_id)) + except Exception as exc: + logger.error(f"Failed to load model {model_id} before deletion: {exc}") + + if not model: + remove_errors.append(f"Model id {model_id} not found in database.") + continue + + try: + deleted = self.db_manager.delete_model(int(model_id)) + except Exception as exc: + logger.error(f"Failed to delete model {model_id}: {exc}") + remove_errors.append(f"Failed to delete model id {model_id} from DB: {exc}") + continue + + if not deleted: + remove_errors.append(f"Model id {model_id} was not deleted (already removed?).") + continue + + deleted_count += 1 + removed, errors = self._delete_model_artifacts_from_disk(model) + removed_paths.extend(removed) + remove_errors.extend(errors) + + # Refresh tabs to reflect the deletion(s). + try: + if hasattr(self, "detection_tab"): + self.detection_tab.refresh() + if hasattr(self, "results_tab"): + self.results_tab.refresh() + if hasattr(self, "validation_tab"): + self.validation_tab.refresh() + if hasattr(self, "training_tab"): + self.training_tab.refresh() + except Exception as exc: + logger.warning(f"Failed to refresh tabs after model deletion: {exc}") + + details: list[str] = [] + if removed_paths: + details.append("Removed from disk:\n" + "\n".join(removed_paths)) + if remove_errors: + details.append("\nDisk cleanup warnings:\n" + "\n".join(remove_errors)) + + QMessageBox.information( + self, + "Delete Model", + f"Deleted {deleted_count} model(s) from database." + ("\n\n" + "\n".join(details) if details else ""), + ) + + def _delete_model(self, model_id: int) -> None: + """Delete a model from the database and remove its artifacts from disk.""" + + model = None + try: + model = self.db_manager.get_model_by_id(model_id) + except Exception as exc: + logger.error(f"Failed to load model {model_id} before deletion: {exc}") + + if not model: + QMessageBox.warning(self, "Delete Model", "Selected model was not found in the database.") + return + + model_path = str(model.get("model_path") or "") + + try: + deleted = self.db_manager.delete_model(model_id) + except Exception as exc: + logger.error(f"Failed to delete model {model_id}: {exc}") + QMessageBox.critical(self, "Delete Model", f"Failed to delete model from database:\n{exc}") + return + + if not deleted: + QMessageBox.warning(self, "Delete Model", "No model was deleted (it may have already been removed).") + return + + removed_paths, remove_errors = self._delete_model_artifacts_from_disk(model) + + # Refresh tabs to reflect the deletion. + try: + if hasattr(self, "detection_tab"): + self.detection_tab.refresh() + if hasattr(self, "results_tab"): + self.results_tab.refresh() + if hasattr(self, "validation_tab"): + self.validation_tab.refresh() + if hasattr(self, "training_tab"): + self.training_tab.refresh() + except Exception as exc: + logger.warning(f"Failed to refresh tabs after model deletion: {exc}") + + details = [] + if model_path: + details.append(f"Deleted model record for: {model_path}") + if removed_paths: + details.append("\nRemoved from disk:\n" + "\n".join(removed_paths)) + if remove_errors: + details.append("\nDisk cleanup warnings:\n" + "\n".join(remove_errors)) + + QMessageBox.information( + self, + "Delete Model", + "Model deleted from database." + ("\n\n" + "\n".join(details) if details else ""), + ) + + def _delete_model_artifacts_from_disk(self, model: dict) -> tuple[list[str], list[str]]: + """Best-effort removal of model artifacts on disk. + + Strategy: + - Remove run directories inferred from: + - model.model_path (…//weights/*.pt => ) + - training_params.stage_results[].results.save_dir + but only if they are under the configured models directory. + - If the weights file itself exists and is outside the models directory, delete only the file. + + Returns: + (removed_paths, errors) + """ + + removed: list[str] = [] + errors: list[str] = [] + + models_root = Path(self.config_manager.get_models_directory() or "data/models").expanduser() + try: + models_root_resolved = models_root.resolve() + except Exception: + models_root_resolved = models_root + + inferred_dirs: list[Path] = [] + + # 1) From model_path + model_path_value = model.get("model_path") + if model_path_value: + try: + p = Path(str(model_path_value)).expanduser() + p_resolved = p.resolve() if p.exists() else p + if p_resolved.is_file(): + if p_resolved.parent.name == "weights" and p_resolved.parent.parent.exists(): + inferred_dirs.append(p_resolved.parent.parent) + elif p_resolved.parent.exists(): + inferred_dirs.append(p_resolved.parent) + except Exception: + pass + + # 2) From training_params.stage_results[].results.save_dir + training_params = model.get("training_params") or {} + if isinstance(training_params, dict): + stage_results = training_params.get("stage_results") + if isinstance(stage_results, list): + for stage in stage_results: + results = (stage or {}).get("results") + save_dir = (results or {}).get("save_dir") if isinstance(results, dict) else None + if not save_dir: + continue + try: + d = Path(str(save_dir)).expanduser() + if d.exists() and d.is_dir(): + inferred_dirs.append(d) + except Exception: + continue + + # Deduplicate inferred_dirs + unique_dirs: list[Path] = [] + seen: set[str] = set() + for d in inferred_dirs: + try: + key = str(d.resolve()) + except Exception: + key = str(d) + if key in seen: + continue + seen.add(key) + unique_dirs.append(d) + + # Delete directories under models_root + for d in unique_dirs: + try: + d_resolved = d.resolve() + except Exception: + d_resolved = d + try: + if d_resolved.exists() and d_resolved.is_dir() and d_resolved.is_relative_to(models_root_resolved): + shutil.rmtree(d_resolved) + removed.append(str(d_resolved)) + except Exception as exc: + errors.append(f"Failed to remove directory {d_resolved}: {exc}") + + # If nothing matched (e.g., model_path outside models_root), delete just the file. + if model_path_value: + try: + p = Path(str(model_path_value)).expanduser() + if p.exists() and p.is_file(): + p_resolved = p.resolve() + if not p_resolved.is_relative_to(models_root_resolved): + p_resolved.unlink() + removed.append(str(p_resolved)) + except Exception as exc: + errors.append(f"Failed to remove model file {model_path_value}: {exc}") + + return removed, errors def _show_about(self): """Show about dialog."""