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

Adding python files

Browse Source
This commit is contained in:
Martin Laasmaa
2025-12-05 09:50:50 +02:00
parent c6143cd11a
commit 6bd2b100ca
24 changed files with 3076 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

344
src/gui/tabs/detection_tab.py Normal file
View File

@@ -0,0 +1,344 @@
"""
Detection tab for the microscopy object detection application.
Handles single image and batch detection.
"""
from PySide6.QtWidgets import (
QWidget,
QVBoxLayout,
QHBoxLayout,
QPushButton,
QLabel,
QComboBox,
QSlider,
QFileDialog,
QMessageBox,
QProgressBar,
QTextEdit,
QGroupBox,
QFormLayout,
)
from PySide6.QtCore import Qt, QThread, Signal
from pathlib import Path
from src.database.db_manager import DatabaseManager
from src.utils.config_manager import ConfigManager
from src.utils.logger import get_logger
from src.utils.file_utils import get_image_files
from src.model.inference import InferenceEngine
logger = get_logger(__name__)
class DetectionWorker(QThread):
"""Worker thread for running detection."""
progress = Signal(int, int, str) # current, total, message
finished = Signal(list) # results
error = Signal(str) # error message
def __init__(self, engine, image_paths, repo_root, conf):
super().__init__()
self.engine = engine
self.image_paths = image_paths
self.repo_root = repo_root
self.conf = conf
def run(self):
"""Run detection in background thread."""
try:
results = self.engine.detect_batch(
self.image_paths, self.repo_root, self.conf, self.progress.emit
)
self.finished.emit(results)
except Exception as e:
logger.error(f"Detection error: {e}")
self.error.emit(str(e))
class DetectionTab(QWidget):
"""Detection tab for single image and batch detection."""
def __init__(
self, db_manager: DatabaseManager, config_manager: ConfigManager, parent=None
):
super().__init__(parent)
self.db_manager = db_manager
self.config_manager = config_manager
self.inference_engine = None
self.current_model_id = None
self._setup_ui()
self._connect_signals()
self._load_models()
def _setup_ui(self):
"""Setup user interface."""
layout = QVBoxLayout()
# Model selection group
model_group = QGroupBox("Model Selection")
model_layout = QFormLayout()
self.model_combo = QComboBox()
self.model_combo.addItem("No models available", None)
model_layout.addRow("Model:", self.model_combo)
model_group.setLayout(model_layout)
layout.addWidget(model_group)
# Detection settings group
settings_group = QGroupBox("Detection Settings")
settings_layout = QFormLayout()
# Confidence threshold
conf_layout = QHBoxLayout()
self.conf_slider = QSlider(Qt.Horizontal)
self.conf_slider.setRange(0, 100)
self.conf_slider.setValue(25)
self.conf_slider.setTickPosition(QSlider.TicksBelow)
self.conf_slider.setTickInterval(10)
conf_layout.addWidget(self.conf_slider)
self.conf_label = QLabel("0.25")
conf_layout.addWidget(self.conf_label)
settings_layout.addRow("Confidence:", conf_layout)
settings_group.setLayout(settings_layout)
layout.addWidget(settings_group)
# Action buttons
button_layout = QHBoxLayout()
self.single_image_btn = QPushButton("Detect Single Image")
self.single_image_btn.clicked.connect(self._detect_single_image)
button_layout.addWidget(self.single_image_btn)
self.batch_btn = QPushButton("Detect Batch (Folder)")
self.batch_btn.clicked.connect(self._detect_batch)
button_layout.addWidget(self.batch_btn)
layout.addLayout(button_layout)
# Progress bar
self.progress_bar = QProgressBar()
self.progress_bar.setVisible(False)
layout.addWidget(self.progress_bar)
# Results display
results_group = QGroupBox("Detection Results")
results_layout = QVBoxLayout()
self.results_text = QTextEdit()
self.results_text.setReadOnly(True)
self.results_text.setMaximumHeight(200)
results_layout.addWidget(self.results_text)
results_group.setLayout(results_layout)
layout.addWidget(results_group)
layout.addStretch()
self.setLayout(layout)
def _connect_signals(self):
"""Connect signals and slots."""
self.conf_slider.valueChanged.connect(self._update_confidence_label)
self.model_combo.currentIndexChanged.connect(self._on_model_changed)
def _load_models(self):
"""Load available models from database."""
try:
models = self.db_manager.get_models()
self.model_combo.clear()
if not models:
self.model_combo.addItem("No models available", None)
self._set_buttons_enabled(False)
return
# Add base model option
base_model = self.config_manager.get(
"models.default_base_model", "yolov8s.pt"
)
self.model_combo.addItem(
f"Base Model ({base_model})", {"id": 0, "path": base_model}
)
# Add trained models
for model in models:
display_name = f"{model['model_name']} v{model['model_version']}"
self.model_combo.addItem(display_name, model)
self._set_buttons_enabled(True)
except Exception as e:
logger.error(f"Error loading models: {e}")
QMessageBox.warning(self, "Error", f"Failed to load models:\n{str(e)}")
def _on_model_changed(self, index: int):
"""Handle model selection change."""
model_data = self.model_combo.itemData(index)
if model_data and model_data["id"] != 0:
self.current_model_id = model_data["id"]
else:
self.current_model_id = None
def _update_confidence_label(self, value: int):
"""Update confidence label."""
conf = value / 100.0
self.conf_label.setText(f"{conf:.2f}")
def _detect_single_image(self):
"""Detect objects in a single image."""
# Get image file
repo_path = self.config_manager.get_image_repository_path()
start_dir = repo_path if repo_path else ""
file_path, _ = QFileDialog.getOpenFileName(
self,
"Select Image",
start_dir,
"Images (*.jpg *.jpeg *.png *.tif *.tiff *.bmp)",
)
if not file_path:
return
# Run detection
self._run_detection([file_path])
def _detect_batch(self):
"""Detect objects in batch (folder)."""
# Get folder
repo_path = self.config_manager.get_image_repository_path()
start_dir = repo_path if repo_path else ""
folder_path = QFileDialog.getExistingDirectory(self, "Select Folder", start_dir)
if not folder_path:
return
# Get all image files
allowed_ext = self.config_manager.get_allowed_extensions()
image_files = get_image_files(folder_path, allowed_ext, recursive=False)
if not image_files:
QMessageBox.information(
self, "No Images", "No image files found in selected folder."
)
return
# Confirm batch processing
reply = QMessageBox.question(
self,
"Confirm Batch Detection",
f"Process {len(image_files)} images?",
QMessageBox.Yes | QMessageBox.No,
)
if reply == QMessageBox.Yes:
self._run_detection(image_files)
def _run_detection(self, image_paths: list):
"""Run detection on image list."""
try:
# Get selected model
model_data = self.model_combo.currentData()
if not model_data:
QMessageBox.warning(self, "No Model", "Please select a model first.")
return
model_path = model_data["path"]
model_id = model_data["id"]
# Ensure we have a valid model ID (create entry for base model if needed)
if model_id == 0:
# Create database entry for base model
base_model = self.config_manager.get(
"models.default_base_model", "yolov8s.pt"
)
model_id = self.db_manager.add_model(
model_name="Base Model",
model_version="pretrained",
model_path=base_model,
base_model=base_model,
)
# Create inference engine
self.inference_engine = InferenceEngine(
model_path, self.db_manager, model_id
)
# Get confidence threshold
conf = self.conf_slider.value() / 100.0
# Get repository root
repo_root = self.config_manager.get_image_repository_path()
if not repo_root:
repo_root = str(Path(image_paths[0]).parent)
# Show progress bar
self.progress_bar.setVisible(True)
self.progress_bar.setMaximum(len(image_paths))
self._set_buttons_enabled(False)
# Create and start worker thread
self.worker = DetectionWorker(
self.inference_engine, image_paths, repo_root, conf
)
self.worker.progress.connect(self._on_progress)
self.worker.finished.connect(self._on_detection_finished)
self.worker.error.connect(self._on_detection_error)
self.worker.start()
except Exception as e:
logger.error(f"Error starting detection: {e}")
QMessageBox.critical(self, "Error", f"Failed to start detection:\n{str(e)}")
self._set_buttons_enabled(True)
def _on_progress(self, current: int, total: int, message: str):
"""Handle progress update."""
self.progress_bar.setValue(current)
self.results_text.append(f"[{current}/{total}] {message}")
def _on_detection_finished(self, results: list):
"""Handle detection completion."""
self.progress_bar.setVisible(False)
self._set_buttons_enabled(True)
# Calculate statistics
total_detections = sum(r["count"] for r in results)
successful = sum(1 for r in results if r.get("success", False))
summary = f"\n=== Detection Complete ===\n"
summary += f"Processed: {len(results)} images\n"
summary += f"Successful: {successful}\n"
summary += f"Total detections: {total_detections}\n"
self.results_text.append(summary)
QMessageBox.information(
self,
"Detection Complete",
f"Processed {len(results)} images\n{total_detections} objects detected",
)
def _on_detection_error(self, error_msg: str):
"""Handle detection error."""
self.progress_bar.setVisible(False)
self._set_buttons_enabled(True)
self.results_text.append(f"\nERROR: {error_msg}")
QMessageBox.critical(self, "Detection Error", error_msg)
def _set_buttons_enabled(self, enabled: bool):
"""Enable/disable action buttons."""
self.single_image_btn.setEnabled(enabled)
self.batch_btn.setEnabled(enabled)
self.model_combo.setEnabled(enabled)
def refresh(self):
"""Refresh the tab."""
self._load_models()
self.results_text.clear()