Adding model deletion feature from database
This commit is contained in:
@@ -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 (…/<run>/weights/*.pt => <run>)
|
||||
- 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."""
|
||||
|
||||
Reference in New Issue
Block a user