Compare commits
40 Commits
c7e1271193
...
monkey-pat
| Author | SHA1 | Date | |
|---|---|---|---|
| d03ffdc4d0 | |||
| 8d30e6bb7a | |||
| f810fec4d8 | |||
| 9c8931e6f3 | |||
| 20578c1fdf | |||
| 2c494dac49 | |||
| 506c74e53a | |||
| eefda5b878 | |||
| 31cb6a6c8e | |||
| 0c19ea2557 | |||
| 89e47591db | |||
| 69cde09e53 | |||
| fcbd5fb16d | |||
| ca52312925 | |||
| 0a93bf797a | |||
| d998c65665 | |||
| 510eabfa94 | |||
| 395d263900 | |||
| e98d287b8a | |||
| d25101de2d | |||
| f88beef188 | |||
| 2fd9a2acf4 | |||
| 2bcd18cc75 | |||
| 5d25378c46 | |||
| 2b0b48921e | |||
| b0c05f0225 | |||
| 97badaa390 | |||
| 8f8132ce61 | |||
| 6ae7481e25 | |||
| 061f8b3ca2 | |||
| a8e5db3135 | |||
| 268ed5175e | |||
| 5e9d3b1dc4 | |||
| 7d83e9b9b1 | |||
| e364d06217 | |||
| e5036c10cf | |||
| c7e388d9ae | |||
| 6b995e7325 | |||
| 0e0741d323 | |||
| dd99a0677c |
@@ -1,57 +0,0 @@
|
||||
database:
|
||||
path: data/detections.db
|
||||
image_repository:
|
||||
base_path: ''
|
||||
allowed_extensions:
|
||||
- .jpg
|
||||
- .jpeg
|
||||
- .png
|
||||
- .tif
|
||||
- .tiff
|
||||
- .bmp
|
||||
models:
|
||||
default_base_model: yolov8s-seg.pt
|
||||
models_directory: data/models
|
||||
base_model_choices:
|
||||
- yolov8s-seg.pt
|
||||
- yolo11s-seg.pt
|
||||
training:
|
||||
default_epochs: 100
|
||||
default_batch_size: 16
|
||||
default_imgsz: 1024
|
||||
default_patience: 50
|
||||
default_lr0: 0.01
|
||||
two_stage:
|
||||
enabled: false
|
||||
stage1:
|
||||
epochs: 20
|
||||
lr0: 0.0005
|
||||
patience: 10
|
||||
freeze: 10
|
||||
stage2:
|
||||
epochs: 150
|
||||
lr0: 0.0003
|
||||
patience: 30
|
||||
last_dataset_yaml: /home/martin/code/object_detection/data/datasets/data.yaml
|
||||
last_dataset_dir: /home/martin/code/object_detection/data/datasets
|
||||
detection:
|
||||
default_confidence: 0.25
|
||||
default_iou: 0.45
|
||||
max_batch_size: 100
|
||||
visualization:
|
||||
bbox_colors:
|
||||
organelle: '#FF6B6B'
|
||||
membrane_branch: '#4ECDC4'
|
||||
default: '#00FF00'
|
||||
bbox_thickness: 2
|
||||
font_size: 12
|
||||
export:
|
||||
formats:
|
||||
- csv
|
||||
- json
|
||||
- excel
|
||||
default_format: csv
|
||||
logging:
|
||||
level: INFO
|
||||
file: logs/app.log
|
||||
format: '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
@@ -82,12 +82,12 @@ include-package-data = true
|
||||
"src.database" = ["*.sql"]
|
||||
|
||||
[tool.black]
|
||||
line-length = 88
|
||||
line-length = 120
|
||||
target-version = ['py38', 'py39', 'py310', 'py311']
|
||||
include = '\.pyi?$'
|
||||
|
||||
[tool.pylint.messages_control]
|
||||
max-line-length = 88
|
||||
max-line-length = 120
|
||||
|
||||
[tool.mypy]
|
||||
python_version = "3.8"
|
||||
|
||||
@@ -60,9 +60,7 @@ class DatabaseManager:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Check if annotations table exists
|
||||
cursor.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='annotations'"
|
||||
)
|
||||
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
|
||||
@@ -203,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(
|
||||
@@ -242,9 +262,7 @@ class DatabaseManager:
|
||||
return cursor.lastrowid
|
||||
except sqlite3.IntegrityError:
|
||||
# Image already exists, return its ID
|
||||
cursor.execute(
|
||||
"SELECT id FROM images WHERE relative_path = ?", (relative_path,)
|
||||
)
|
||||
cursor.execute("SELECT id FROM images WHERE relative_path = ?", (relative_path,))
|
||||
row = cursor.fetchone()
|
||||
return row["id"] if row else None
|
||||
finally:
|
||||
@@ -255,17 +273,13 @@ class DatabaseManager:
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT * FROM images WHERE relative_path = ?", (relative_path,)
|
||||
)
|
||||
cursor.execute("SELECT * FROM images WHERE relative_path = ?", (relative_path,))
|
||||
row = cursor.fetchone()
|
||||
return dict(row) if row else None
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def get_or_create_image(
|
||||
self, relative_path: str, filename: str, width: int, height: int
|
||||
) -> int:
|
||||
def get_or_create_image(self, relative_path: str, filename: str, width: int, height: int) -> int:
|
||||
"""Get existing image or create new one."""
|
||||
existing = self.get_image_by_path(relative_path)
|
||||
if existing:
|
||||
@@ -355,16 +369,8 @@ class DatabaseManager:
|
||||
bbox[2],
|
||||
bbox[3],
|
||||
det["confidence"],
|
||||
(
|
||||
json.dumps(det.get("segmentation_mask"))
|
||||
if det.get("segmentation_mask")
|
||||
else None
|
||||
),
|
||||
(
|
||||
json.dumps(det.get("metadata"))
|
||||
if det.get("metadata")
|
||||
else None
|
||||
),
|
||||
(json.dumps(det.get("segmentation_mask")) if det.get("segmentation_mask") else None),
|
||||
(json.dumps(det.get("metadata")) if det.get("metadata") else None),
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
@@ -409,12 +415,13 @@ class DatabaseManager:
|
||||
if filters:
|
||||
conditions = []
|
||||
for key, value in filters.items():
|
||||
if (
|
||||
key.startswith("d.")
|
||||
or key.startswith("i.")
|
||||
or key.startswith("m.")
|
||||
):
|
||||
if key.startswith("d.") or key.startswith("i.") or key.startswith("m."):
|
||||
if "like" in value.lower():
|
||||
conditions.append(f"{key} LIKE ?")
|
||||
params.append(value.split(" ")[1])
|
||||
else:
|
||||
conditions.append(f"{key} = ?")
|
||||
params.append(value)
|
||||
else:
|
||||
conditions.append(f"d.{key} = ?")
|
||||
params.append(value)
|
||||
@@ -442,18 +449,14 @@ class DatabaseManager:
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def get_detections_for_image(
|
||||
self, image_id: int, model_id: Optional[int] = None
|
||||
) -> List[Dict]:
|
||||
def get_detections_for_image(self, image_id: int, model_id: Optional[int] = None) -> List[Dict]:
|
||||
"""Get all detections for a specific image."""
|
||||
filters = {"image_id": image_id}
|
||||
if model_id:
|
||||
filters["model_id"] = model_id
|
||||
return self.get_detections(filters)
|
||||
|
||||
def delete_detections_for_image(
|
||||
self, image_id: int, model_id: Optional[int] = None
|
||||
) -> int:
|
||||
def delete_detections_for_image(self, image_id: int, model_id: Optional[int] = None) -> int:
|
||||
"""Delete detections tied to a specific image and optional model."""
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
@@ -481,6 +484,22 @@ class DatabaseManager:
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def delete_all_detections(self) -> int:
|
||||
"""Delete all detections from the database.
|
||||
|
||||
Returns:
|
||||
Number of rows deleted.
|
||||
"""
|
||||
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("DELETE FROM detections")
|
||||
conn.commit()
|
||||
return cursor.rowcount
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
# ==================== Statistics Operations ====================
|
||||
|
||||
def get_detection_statistics(
|
||||
@@ -524,9 +543,7 @@ class DatabaseManager:
|
||||
""",
|
||||
params,
|
||||
)
|
||||
class_counts = {
|
||||
row["class_name"]: row["count"] for row in cursor.fetchall()
|
||||
}
|
||||
class_counts = {row["class_name"]: row["count"] for row in cursor.fetchall()}
|
||||
|
||||
# Average confidence
|
||||
cursor.execute(
|
||||
@@ -583,9 +600,7 @@ class DatabaseManager:
|
||||
|
||||
# ==================== Export Operations ====================
|
||||
|
||||
def export_detections_to_csv(
|
||||
self, output_path: str, filters: Optional[Dict] = None
|
||||
) -> bool:
|
||||
def export_detections_to_csv(self, output_path: str, filters: Optional[Dict] = None) -> bool:
|
||||
"""Export detections to CSV file."""
|
||||
try:
|
||||
detections = self.get_detections(filters)
|
||||
@@ -614,9 +629,7 @@ class DatabaseManager:
|
||||
for det in detections:
|
||||
row = {k: det[k] for k in fieldnames if k in det}
|
||||
# Convert segmentation mask list to JSON string for CSV
|
||||
if row.get("segmentation_mask") and isinstance(
|
||||
row["segmentation_mask"], list
|
||||
):
|
||||
if row.get("segmentation_mask") and isinstance(row["segmentation_mask"], list):
|
||||
row["segmentation_mask"] = json.dumps(row["segmentation_mask"])
|
||||
writer.writerow(row)
|
||||
|
||||
@@ -625,9 +638,7 @@ class DatabaseManager:
|
||||
print(f"Error exporting to CSV: {e}")
|
||||
return False
|
||||
|
||||
def export_detections_to_json(
|
||||
self, output_path: str, filters: Optional[Dict] = None
|
||||
) -> bool:
|
||||
def export_detections_to_json(self, output_path: str, filters: Optional[Dict] = None) -> bool:
|
||||
"""Export detections to JSON file."""
|
||||
try:
|
||||
detections = self.get_detections(filters)
|
||||
@@ -647,6 +658,75 @@ class DatabaseManager:
|
||||
|
||||
# ==================== Annotation Operations ====================
|
||||
|
||||
def get_annotated_images_summary(
|
||||
self,
|
||||
name_filter: Optional[str] = None,
|
||||
order_by: str = "filename",
|
||||
order_dir: str = "ASC",
|
||||
limit: Optional[int] = None,
|
||||
offset: int = 0,
|
||||
) -> List[Dict]:
|
||||
"""Return images that have at least one manual annotation.
|
||||
|
||||
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_sql = ""
|
||||
if name_filter:
|
||||
# Case-insensitive substring search.
|
||||
token = f"%{name_filter}%"
|
||||
where_sql = "WHERE (i.filename LIKE ? OR i.relative_path LIKE ?)"
|
||||
params.extend([token, token])
|
||||
|
||||
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
|
||||
JOIN annotations a ON a.image_id = i.id
|
||||
{where_sql}
|
||||
GROUP BY i.id
|
||||
HAVING annotation_count > 0
|
||||
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 add_annotation(
|
||||
self,
|
||||
image_id: int,
|
||||
@@ -785,17 +865,13 @@ class DatabaseManager:
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT * FROM object_classes WHERE class_name = ?", (class_name,)
|
||||
)
|
||||
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:
|
||||
def add_object_class(self, class_name: str, color: str, description: Optional[str] = None) -> int:
|
||||
"""
|
||||
Add a new object class.
|
||||
|
||||
@@ -928,8 +1004,7 @@ class DatabaseManager:
|
||||
if not split_map[required]:
|
||||
raise ValueError(
|
||||
"Unable to determine %s image directory under %s. Provide it "
|
||||
"explicitly via the 'splits' argument."
|
||||
% (required, dataset_root_path)
|
||||
"explicitly via the 'splits' argument." % (required, dataset_root_path)
|
||||
)
|
||||
|
||||
yaml_splits: Dict[str, str] = {}
|
||||
@@ -955,11 +1030,7 @@ class DatabaseManager:
|
||||
if yaml_splits.get("test"):
|
||||
payload["test"] = yaml_splits["test"]
|
||||
|
||||
output_path_obj = (
|
||||
Path(output_path).expanduser()
|
||||
if output_path
|
||||
else dataset_root_path / "data.yaml"
|
||||
)
|
||||
output_path_obj = Path(output_path).expanduser() if output_path else dataset_root_path / "data.yaml"
|
||||
output_path_obj.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(output_path_obj, "w", encoding="utf-8") as handle:
|
||||
@@ -1019,15 +1090,9 @@ class DatabaseManager:
|
||||
for split_name, options in patterns.items():
|
||||
for relative in options:
|
||||
candidate = (dataset_root / relative).resolve()
|
||||
if (
|
||||
candidate.exists()
|
||||
and candidate.is_dir()
|
||||
and self._directory_has_images(candidate)
|
||||
):
|
||||
if candidate.exists() and candidate.is_dir() and self._directory_has_images(candidate):
|
||||
try:
|
||||
inferred[split_name] = candidate.relative_to(
|
||||
dataset_root
|
||||
).as_posix()
|
||||
inferred[split_name] = candidate.relative_to(dataset_root).as_posix()
|
||||
except ValueError:
|
||||
inferred[split_name] = candidate.as_posix()
|
||||
break
|
||||
|
||||
@@ -55,10 +55,7 @@ CREATE TABLE IF NOT EXISTS object_classes (
|
||||
|
||||
-- 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');
|
||||
('terminal', '#FFFF00', 'Axion terminal');
|
||||
|
||||
-- Annotations table: stores manual annotations
|
||||
CREATE TABLE IF NOT EXISTS annotations (
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -117,10 +125,10 @@ class MainWindow(QMainWindow):
|
||||
|
||||
# Add tabs to widget
|
||||
self.tab_widget.addTab(self.detection_tab, "Detection")
|
||||
self.tab_widget.addTab(self.results_tab, "Results")
|
||||
self.tab_widget.addTab(self.annotation_tab, "Annotation")
|
||||
self.tab_widget.addTab(self.training_tab, "Training")
|
||||
self.tab_widget.addTab(self.validation_tab, "Validation")
|
||||
self.tab_widget.addTab(self.results_tab, "Results")
|
||||
self.tab_widget.addTab(self.annotation_tab, "Annotation (Future)")
|
||||
|
||||
# Connect tab change signal
|
||||
self.tab_widget.currentChanged.connect(self._on_tab_changed)
|
||||
@@ -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."""
|
||||
@@ -193,6 +199,10 @@ class MainWindow(QMainWindow):
|
||||
self.training_tab.refresh()
|
||||
if hasattr(self, "results_tab"):
|
||||
self.results_tab.refresh()
|
||||
if hasattr(self, "annotation_tab"):
|
||||
self.annotation_tab.refresh()
|
||||
if hasattr(self, "validation_tab"):
|
||||
self.validation_tab.refresh()
|
||||
except Exception as e:
|
||||
logger.error(f"Error applying settings: {e}")
|
||||
|
||||
@@ -209,6 +219,14 @@ class MainWindow(QMainWindow):
|
||||
logger.debug(f"Switched to tab: {tab_name}")
|
||||
self._update_status(f"Viewing: {tab_name}")
|
||||
|
||||
# Ensure the Annotation tab always shows up-to-date DB-backed lists.
|
||||
try:
|
||||
current_widget = self.tab_widget.widget(index)
|
||||
if hasattr(self, "annotation_tab") and current_widget is self.annotation_tab:
|
||||
self.annotation_tab.refresh()
|
||||
except Exception as exc:
|
||||
logger.debug(f"Failed to refresh annotation tab on selection: {exc}")
|
||||
|
||||
def _show_database_stats(self):
|
||||
"""Show database statistics dialog."""
|
||||
try:
|
||||
@@ -231,10 +249,230 @@ 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."""
|
||||
about_text = """
|
||||
@@ -301,6 +539,11 @@ class MainWindow(QMainWindow):
|
||||
if hasattr(self, "training_tab"):
|
||||
self.training_tab.shutdown()
|
||||
if hasattr(self, "annotation_tab"):
|
||||
# Best-effort refresh so DB-backed UI state is consistent at shutdown.
|
||||
try:
|
||||
self.annotation_tab.refresh()
|
||||
except Exception:
|
||||
pass
|
||||
self.annotation_tab.save_state()
|
||||
|
||||
logger.info("Application closing")
|
||||
|
||||
@@ -13,6 +13,11 @@ from PySide6.QtWidgets import (
|
||||
QFileDialog,
|
||||
QMessageBox,
|
||||
QSplitter,
|
||||
QLineEdit,
|
||||
QTableWidget,
|
||||
QTableWidgetItem,
|
||||
QHeaderView,
|
||||
QAbstractItemView,
|
||||
)
|
||||
from PySide6.QtCore import Qt, QSettings
|
||||
from pathlib import Path
|
||||
@@ -29,9 +34,7 @@ logger = get_logger(__name__)
|
||||
class AnnotationTab(QWidget):
|
||||
"""Annotation tab for manual image annotation."""
|
||||
|
||||
def __init__(
|
||||
self, db_manager: DatabaseManager, config_manager: ConfigManager, parent=None
|
||||
):
|
||||
def __init__(self, db_manager: DatabaseManager, config_manager: ConfigManager, parent=None):
|
||||
super().__init__(parent)
|
||||
self.db_manager = db_manager
|
||||
self.config_manager = config_manager
|
||||
@@ -52,6 +55,32 @@ class AnnotationTab(QWidget):
|
||||
self.main_splitter = QSplitter(Qt.Horizontal)
|
||||
self.main_splitter.setHandleWidth(10)
|
||||
|
||||
# { Left-most pane: annotated images list
|
||||
annotated_group = QGroupBox("Annotated Images")
|
||||
annotated_layout = QVBoxLayout()
|
||||
|
||||
filter_row = QHBoxLayout()
|
||||
filter_row.addWidget(QLabel("Filter:"))
|
||||
self.annotated_filter_edit = QLineEdit()
|
||||
self.annotated_filter_edit.setPlaceholderText("Type to filter by image name…")
|
||||
self.annotated_filter_edit.textChanged.connect(self._refresh_annotated_images_list)
|
||||
filter_row.addWidget(self.annotated_filter_edit, 1)
|
||||
annotated_layout.addLayout(filter_row)
|
||||
|
||||
self.annotated_images_table = QTableWidget(0, 2)
|
||||
self.annotated_images_table.setHorizontalHeaderLabels(["Image", "Annotations"])
|
||||
self.annotated_images_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.Stretch)
|
||||
self.annotated_images_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeToContents)
|
||||
self.annotated_images_table.setSelectionBehavior(QAbstractItemView.SelectRows)
|
||||
self.annotated_images_table.setSelectionMode(QAbstractItemView.SingleSelection)
|
||||
self.annotated_images_table.setEditTriggers(QAbstractItemView.NoEditTriggers)
|
||||
self.annotated_images_table.setSortingEnabled(True)
|
||||
self.annotated_images_table.itemSelectionChanged.connect(self._on_annotated_image_selected)
|
||||
annotated_layout.addWidget(self.annotated_images_table, 1)
|
||||
|
||||
annotated_group.setLayout(annotated_layout)
|
||||
# }
|
||||
|
||||
# { Left splitter for image display and zoom info
|
||||
self.left_splitter = QSplitter(Qt.Vertical)
|
||||
self.left_splitter.setHandleWidth(10)
|
||||
@@ -62,6 +91,9 @@ class AnnotationTab(QWidget):
|
||||
|
||||
# Use the AnnotationCanvasWidget
|
||||
self.annotation_canvas = AnnotationCanvasWidget()
|
||||
# Auto-zoom so newly loaded images fill the available canvas viewport.
|
||||
# (Matches the behavior used in ResultsTab.)
|
||||
self.annotation_canvas.set_auto_fit_to_view(True)
|
||||
self.annotation_canvas.zoom_changed.connect(self._on_zoom_changed)
|
||||
self.annotation_canvas.annotation_drawn.connect(self._on_annotation_drawn)
|
||||
# Selection of existing polylines (when tool is not in drawing mode)
|
||||
@@ -72,9 +104,7 @@ class AnnotationTab(QWidget):
|
||||
self.left_splitter.addWidget(canvas_group)
|
||||
|
||||
# Controls info
|
||||
controls_info = QLabel(
|
||||
"Zoom: Mouse wheel or +/- keys | Drawing: Enable pen and drag mouse"
|
||||
)
|
||||
controls_info = QLabel("Zoom: Mouse wheel or +/- keys | Drawing: Enable pen and drag mouse")
|
||||
controls_info.setStyleSheet("QLabel { color: #888; font-style: italic; }")
|
||||
self.left_splitter.addWidget(controls_info)
|
||||
# }
|
||||
@@ -85,36 +115,20 @@ class AnnotationTab(QWidget):
|
||||
|
||||
# Annotation tools section
|
||||
self.annotation_tools = AnnotationToolsWidget(self.db_manager)
|
||||
self.annotation_tools.polyline_enabled_changed.connect(
|
||||
self.annotation_canvas.set_polyline_enabled
|
||||
)
|
||||
self.annotation_tools.polyline_pen_color_changed.connect(
|
||||
self.annotation_canvas.set_polyline_pen_color
|
||||
)
|
||||
self.annotation_tools.polyline_pen_width_changed.connect(
|
||||
self.annotation_canvas.set_polyline_pen_width
|
||||
)
|
||||
self.annotation_tools.polyline_enabled_changed.connect(self.annotation_canvas.set_polyline_enabled)
|
||||
self.annotation_tools.polyline_pen_color_changed.connect(self.annotation_canvas.set_polyline_pen_color)
|
||||
self.annotation_tools.polyline_pen_width_changed.connect(self.annotation_canvas.set_polyline_pen_width)
|
||||
# Show / hide bounding boxes
|
||||
self.annotation_tools.show_bboxes_changed.connect(
|
||||
self.annotation_canvas.set_show_bboxes
|
||||
)
|
||||
self.annotation_tools.show_bboxes_changed.connect(self.annotation_canvas.set_show_bboxes)
|
||||
# RDP simplification controls
|
||||
self.annotation_tools.simplify_on_finish_changed.connect(
|
||||
self._on_simplify_on_finish_changed
|
||||
)
|
||||
self.annotation_tools.simplify_epsilon_changed.connect(
|
||||
self._on_simplify_epsilon_changed
|
||||
)
|
||||
self.annotation_tools.simplify_on_finish_changed.connect(self._on_simplify_on_finish_changed)
|
||||
self.annotation_tools.simplify_epsilon_changed.connect(self._on_simplify_epsilon_changed)
|
||||
# Class selection and class-color changes
|
||||
self.annotation_tools.class_selected.connect(self._on_class_selected)
|
||||
self.annotation_tools.class_color_changed.connect(self._on_class_color_changed)
|
||||
self.annotation_tools.clear_annotations_requested.connect(
|
||||
self._on_clear_annotations
|
||||
)
|
||||
self.annotation_tools.clear_annotations_requested.connect(self._on_clear_annotations)
|
||||
# Delete selected annotation on canvas
|
||||
self.annotation_tools.delete_selected_annotation_requested.connect(
|
||||
self._on_delete_selected_annotation
|
||||
)
|
||||
self.annotation_tools.delete_selected_annotation_requested.connect(self._on_delete_selected_annotation)
|
||||
self.right_splitter.addWidget(self.annotation_tools)
|
||||
|
||||
# Image loading section
|
||||
@@ -137,12 +151,13 @@ class AnnotationTab(QWidget):
|
||||
self.right_splitter.addWidget(load_group)
|
||||
# }
|
||||
|
||||
# Add both splitters to the main horizontal splitter
|
||||
# Add list + both splitters to the main horizontal splitter
|
||||
self.main_splitter.addWidget(annotated_group)
|
||||
self.main_splitter.addWidget(self.left_splitter)
|
||||
self.main_splitter.addWidget(self.right_splitter)
|
||||
|
||||
# Set initial sizes: 75% for left (image), 25% for right (controls)
|
||||
self.main_splitter.setSizes([750, 250])
|
||||
# Set initial sizes: list (left), canvas (middle), controls (right)
|
||||
self.main_splitter.setSizes([320, 650, 280])
|
||||
|
||||
layout.addWidget(self.main_splitter)
|
||||
self.setLayout(layout)
|
||||
@@ -150,6 +165,9 @@ class AnnotationTab(QWidget):
|
||||
# Restore splitter positions from settings
|
||||
self._restore_state()
|
||||
|
||||
# Populate list on startup.
|
||||
self._refresh_annotated_images_list()
|
||||
|
||||
def _load_image(self):
|
||||
"""Load and display an image file."""
|
||||
# Get last opened directory from QSettings
|
||||
@@ -180,12 +198,24 @@ class AnnotationTab(QWidget):
|
||||
self.current_image_path = file_path
|
||||
|
||||
# Store the directory for next time
|
||||
settings.setValue(
|
||||
"annotation_tab/last_directory", str(Path(file_path).parent)
|
||||
)
|
||||
settings.setValue("annotation_tab/last_directory", str(Path(file_path).parent))
|
||||
|
||||
# Get or create image in database
|
||||
relative_path = str(Path(file_path).name) # Simplified for now
|
||||
repo_root = self.config_manager.get_image_repository_path()
|
||||
relative_path: str
|
||||
try:
|
||||
if repo_root:
|
||||
repo_root_path = Path(repo_root).expanduser().resolve()
|
||||
file_resolved = Path(file_path).expanduser().resolve()
|
||||
if file_resolved.is_relative_to(repo_root_path):
|
||||
relative_path = file_resolved.relative_to(repo_root_path).as_posix()
|
||||
else:
|
||||
# Fallback: store filename only to avoid leaking absolute paths.
|
||||
relative_path = file_resolved.name
|
||||
else:
|
||||
relative_path = str(Path(file_path).name)
|
||||
except Exception:
|
||||
relative_path = str(Path(file_path).name)
|
||||
self.current_image_id = self.db_manager.get_or_create_image(
|
||||
relative_path,
|
||||
Path(file_path).name,
|
||||
@@ -199,6 +229,9 @@ class AnnotationTab(QWidget):
|
||||
# Load and display any existing annotations for this image
|
||||
self._load_annotations_for_current_image()
|
||||
|
||||
# Update annotated images list (newly annotated image added/selected).
|
||||
self._refresh_annotated_images_list(select_image_id=self.current_image_id)
|
||||
|
||||
# Update info label
|
||||
self._update_image_info()
|
||||
|
||||
@@ -206,9 +239,7 @@ class AnnotationTab(QWidget):
|
||||
|
||||
except ImageLoadError as e:
|
||||
logger.error(f"Failed to load image: {e}")
|
||||
QMessageBox.critical(
|
||||
self, "Error Loading Image", f"Failed to load image:\n{str(e)}"
|
||||
)
|
||||
QMessageBox.critical(self, "Error Loading Image", f"Failed to load image:\n{str(e)}")
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error loading image: {e}")
|
||||
QMessageBox.critical(self, "Error", f"Unexpected error:\n{str(e)}")
|
||||
@@ -296,6 +327,9 @@ class AnnotationTab(QWidget):
|
||||
# Reload annotations from DB and redraw (respecting current class filter)
|
||||
self._load_annotations_for_current_image()
|
||||
|
||||
# Update list counts.
|
||||
self._refresh_annotated_images_list(select_image_id=self.current_image_id)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to save annotation: {e}")
|
||||
QMessageBox.critical(self, "Error", f"Failed to save annotation:\n{str(e)}")
|
||||
@@ -340,9 +374,7 @@ class AnnotationTab(QWidget):
|
||||
if not self.current_image_id:
|
||||
return
|
||||
|
||||
logger.debug(
|
||||
f"Class color changed; reloading annotations for image ID {self.current_image_id}"
|
||||
)
|
||||
logger.debug(f"Class color changed; reloading annotations for image ID {self.current_image_id}")
|
||||
self._load_annotations_for_current_image()
|
||||
|
||||
def _on_class_selected(self, class_data):
|
||||
@@ -355,9 +387,7 @@ class AnnotationTab(QWidget):
|
||||
if class_data:
|
||||
logger.debug(f"Object class selected: {class_data['class_name']}")
|
||||
else:
|
||||
logger.debug(
|
||||
'No class selected ("-- Select Class --"), showing all annotations'
|
||||
)
|
||||
logger.debug('No class selected ("-- Select Class --"), showing all annotations')
|
||||
|
||||
# Changing the class filter invalidates any previous selection
|
||||
self.selected_annotation_ids = []
|
||||
@@ -390,9 +420,7 @@ class AnnotationTab(QWidget):
|
||||
question = "Are you sure you want to delete the selected annotation?"
|
||||
title = "Delete Annotation"
|
||||
else:
|
||||
question = (
|
||||
f"Are you sure you want to delete the {count} selected annotations?"
|
||||
)
|
||||
question = f"Are you sure you want to delete the {count} selected annotations?"
|
||||
title = "Delete Annotations"
|
||||
|
||||
reply = QMessageBox.question(
|
||||
@@ -420,13 +448,11 @@ class AnnotationTab(QWidget):
|
||||
QMessageBox.warning(
|
||||
self,
|
||||
"Partial Failure",
|
||||
"Some annotations could not be deleted:\n"
|
||||
+ ", ".join(str(a) for a in failed_ids),
|
||||
"Some annotations could not be deleted:\n" + ", ".join(str(a) for a in failed_ids),
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
f"Deleted {count} annotation(s): "
|
||||
+ ", ".join(str(a) for a in self.selected_annotation_ids)
|
||||
f"Deleted {count} annotation(s): " + ", ".join(str(a) for a in self.selected_annotation_ids)
|
||||
)
|
||||
|
||||
# Clear selection and reload annotations for the current image from DB
|
||||
@@ -434,6 +460,9 @@ class AnnotationTab(QWidget):
|
||||
self.annotation_tools.set_has_selected_annotation(False)
|
||||
self._load_annotations_for_current_image()
|
||||
|
||||
# Update list counts.
|
||||
self._refresh_annotated_images_list(select_image_id=self.current_image_id)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete annotations: {e}")
|
||||
QMessageBox.critical(
|
||||
@@ -456,17 +485,13 @@ class AnnotationTab(QWidget):
|
||||
return
|
||||
|
||||
try:
|
||||
self.current_annotations = self.db_manager.get_annotations_for_image(
|
||||
self.current_image_id
|
||||
)
|
||||
self.current_annotations = self.db_manager.get_annotations_for_image(self.current_image_id)
|
||||
# New annotations loaded; reset any selection
|
||||
self.selected_annotation_ids = []
|
||||
self.annotation_tools.set_has_selected_annotation(False)
|
||||
self._redraw_annotations_for_current_filter()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to load annotations for image {self.current_image_id}: {e}"
|
||||
)
|
||||
logger.error(f"Failed to load annotations for image {self.current_image_id}: {e}")
|
||||
QMessageBox.critical(
|
||||
self,
|
||||
"Error",
|
||||
@@ -490,10 +515,7 @@ class AnnotationTab(QWidget):
|
||||
drawn_count = 0
|
||||
for ann in self.current_annotations:
|
||||
# Filter by class if one is selected
|
||||
if (
|
||||
selected_class_id is not None
|
||||
and ann.get("class_id") != selected_class_id
|
||||
):
|
||||
if selected_class_id is not None and ann.get("class_id") != selected_class_id:
|
||||
continue
|
||||
|
||||
if ann.get("segmentation_mask"):
|
||||
@@ -545,22 +567,176 @@ class AnnotationTab(QWidget):
|
||||
settings = QSettings("microscopy_app", "object_detection")
|
||||
|
||||
# Save main splitter state
|
||||
settings.setValue(
|
||||
"annotation_tab/main_splitter_state", self.main_splitter.saveState()
|
||||
)
|
||||
settings.setValue("annotation_tab/main_splitter_state", self.main_splitter.saveState())
|
||||
|
||||
# Save left splitter state
|
||||
settings.setValue(
|
||||
"annotation_tab/left_splitter_state", self.left_splitter.saveState()
|
||||
)
|
||||
settings.setValue("annotation_tab/left_splitter_state", self.left_splitter.saveState())
|
||||
|
||||
# Save right splitter state
|
||||
settings.setValue(
|
||||
"annotation_tab/right_splitter_state", self.right_splitter.saveState()
|
||||
)
|
||||
settings.setValue("annotation_tab/right_splitter_state", self.right_splitter.saveState())
|
||||
|
||||
logger.debug("Saved annotation tab splitter states")
|
||||
|
||||
def refresh(self):
|
||||
"""Refresh the tab."""
|
||||
self._refresh_annotated_images_list(select_image_id=self.current_image_id)
|
||||
|
||||
# ==================== Annotated images list ====================
|
||||
|
||||
def _refresh_annotated_images_list(self, select_image_id: int | None = None) -> None:
|
||||
"""Reload annotated-images list from the database."""
|
||||
if not hasattr(self, "annotated_images_table"):
|
||||
return
|
||||
|
||||
# Preserve selection if possible
|
||||
desired_id = select_image_id if select_image_id is not None else self.current_image_id
|
||||
|
||||
name_filter = ""
|
||||
if hasattr(self, "annotated_filter_edit"):
|
||||
name_filter = self.annotated_filter_edit.text().strip()
|
||||
|
||||
try:
|
||||
rows = self.db_manager.get_annotated_images_summary(name_filter=name_filter)
|
||||
except Exception as exc:
|
||||
logger.error(f"Failed to load annotated images summary: {exc}")
|
||||
rows = []
|
||||
|
||||
sorting_enabled = self.annotated_images_table.isSortingEnabled()
|
||||
self.annotated_images_table.setSortingEnabled(False)
|
||||
self.annotated_images_table.blockSignals(True)
|
||||
try:
|
||||
self.annotated_images_table.setRowCount(len(rows))
|
||||
for r, entry in enumerate(rows):
|
||||
image_name = str(entry.get("filename") or "")
|
||||
count = int(entry.get("annotation_count") or 0)
|
||||
rel_path = str(entry.get("relative_path") or "")
|
||||
|
||||
name_item = QTableWidgetItem(image_name)
|
||||
# Tooltip shows full path of the image (best-effort: repository_root + relative_path)
|
||||
full_path = rel_path
|
||||
repo_root = self.config_manager.get_image_repository_path()
|
||||
if repo_root and rel_path and not Path(rel_path).is_absolute():
|
||||
try:
|
||||
full_path = str((Path(repo_root) / rel_path).resolve())
|
||||
except Exception:
|
||||
full_path = str(Path(repo_root) / rel_path)
|
||||
name_item.setToolTip(full_path)
|
||||
name_item.setData(Qt.UserRole, int(entry.get("id")))
|
||||
name_item.setData(Qt.UserRole + 1, rel_path)
|
||||
|
||||
count_item = QTableWidgetItem()
|
||||
# Use EditRole to ensure numeric sorting.
|
||||
count_item.setData(Qt.EditRole, count)
|
||||
count_item.setData(Qt.UserRole, int(entry.get("id")))
|
||||
count_item.setData(Qt.UserRole + 1, rel_path)
|
||||
|
||||
self.annotated_images_table.setItem(r, 0, name_item)
|
||||
self.annotated_images_table.setItem(r, 1, count_item)
|
||||
|
||||
# Re-select desired row
|
||||
if desired_id is not None:
|
||||
for r in range(self.annotated_images_table.rowCount()):
|
||||
item = self.annotated_images_table.item(r, 0)
|
||||
if item and item.data(Qt.UserRole) == desired_id:
|
||||
self.annotated_images_table.selectRow(r)
|
||||
break
|
||||
finally:
|
||||
self.annotated_images_table.blockSignals(False)
|
||||
self.annotated_images_table.setSortingEnabled(sorting_enabled)
|
||||
|
||||
def _on_annotated_image_selected(self) -> None:
|
||||
"""When user clicks an item in the list, load that image in the annotation canvas."""
|
||||
selected = self.annotated_images_table.selectedItems()
|
||||
if not selected:
|
||||
return
|
||||
|
||||
# Row selection -> take the first column item
|
||||
row = self.annotated_images_table.currentRow()
|
||||
item = self.annotated_images_table.item(row, 0)
|
||||
if not item:
|
||||
return
|
||||
|
||||
image_id = item.data(Qt.UserRole)
|
||||
rel_path = item.data(Qt.UserRole + 1) or ""
|
||||
if not image_id:
|
||||
return
|
||||
|
||||
image_path = self._resolve_image_path_for_relative_path(rel_path)
|
||||
if not image_path:
|
||||
QMessageBox.warning(
|
||||
self,
|
||||
"Image Not Found",
|
||||
"Unable to locate image on disk for:\n"
|
||||
f"{rel_path}\n\n"
|
||||
"Tip: set Settings → Image repository path to the folder containing your images.",
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
self.current_image = Image(image_path)
|
||||
self.current_image_path = image_path
|
||||
self.current_image_id = int(image_id)
|
||||
self.annotation_canvas.load_image(self.current_image)
|
||||
self._load_annotations_for_current_image()
|
||||
self._update_image_info()
|
||||
except ImageLoadError as exc:
|
||||
logger.error(f"Failed to load image '{image_path}': {exc}")
|
||||
QMessageBox.critical(self, "Error Loading Image", f"Failed to load image:\n{exc}")
|
||||
except Exception as exc:
|
||||
logger.error(f"Unexpected error loading image '{image_path}': {exc}")
|
||||
QMessageBox.critical(self, "Error", f"Unexpected error:\n{exc}")
|
||||
|
||||
def _resolve_image_path_for_relative_path(self, relative_path: str) -> str | None:
|
||||
"""Best-effort conversion from a DB relative_path to an on-disk file path."""
|
||||
|
||||
rel = (relative_path or "").strip()
|
||||
if not rel:
|
||||
return None
|
||||
|
||||
candidates: list[Path] = []
|
||||
|
||||
# 1) Repository root + relative
|
||||
repo_root = (self.config_manager.get_image_repository_path() or "").strip()
|
||||
if repo_root:
|
||||
candidates.append(Path(repo_root) / rel)
|
||||
|
||||
# 2) If the DB path is absolute, try it directly.
|
||||
candidates.append(Path(rel))
|
||||
|
||||
# 3) Try the directory of the currently loaded image (helps when DB stores only filenames)
|
||||
if self.current_image_path:
|
||||
try:
|
||||
candidates.append(Path(self.current_image_path).expanduser().resolve().parent / Path(rel).name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 4) Try the last directory used by the annotation file picker
|
||||
try:
|
||||
settings = QSettings("microscopy_app", "object_detection")
|
||||
last_dir = settings.value("annotation_tab/last_directory", None)
|
||||
if last_dir:
|
||||
candidates.append(Path(str(last_dir)) / Path(rel).name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
for p in candidates:
|
||||
try:
|
||||
expanded = p.expanduser()
|
||||
if expanded.exists() and expanded.is_file():
|
||||
return str(expanded.resolve())
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
# 5) Fallback: search by filename within repository root.
|
||||
filename = Path(rel).name
|
||||
if repo_root and filename:
|
||||
root = Path(repo_root).expanduser()
|
||||
try:
|
||||
if root.exists():
|
||||
for match in root.rglob(filename):
|
||||
if match.is_file():
|
||||
return str(match.resolve())
|
||||
except Exception as exc:
|
||||
logger.debug(f"Search for {filename} under {root} failed: {exc}")
|
||||
|
||||
return None
|
||||
|
||||
@@ -3,7 +3,7 @@ Results tab for browsing stored detections and visualizing overlays.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from PySide6.QtWidgets import (
|
||||
QWidget,
|
||||
@@ -35,9 +35,7 @@ logger = get_logger(__name__)
|
||||
class ResultsTab(QWidget):
|
||||
"""Results tab showing detection history and preview overlays."""
|
||||
|
||||
def __init__(
|
||||
self, db_manager: DatabaseManager, config_manager: ConfigManager, parent=None
|
||||
):
|
||||
def __init__(self, db_manager: DatabaseManager, config_manager: ConfigManager, parent=None):
|
||||
super().__init__(parent)
|
||||
self.db_manager = db_manager
|
||||
self.config_manager = config_manager
|
||||
@@ -67,28 +65,32 @@ class ResultsTab(QWidget):
|
||||
self.refresh_btn = QPushButton("Refresh")
|
||||
self.refresh_btn.clicked.connect(self.refresh)
|
||||
controls_layout.addWidget(self.refresh_btn)
|
||||
|
||||
self.delete_all_btn = QPushButton("Delete All Detections")
|
||||
self.delete_all_btn.setToolTip(
|
||||
"Permanently delete ALL detections from the database.\n" "This cannot be undone."
|
||||
)
|
||||
self.delete_all_btn.clicked.connect(self._delete_all_detections)
|
||||
controls_layout.addWidget(self.delete_all_btn)
|
||||
|
||||
self.export_labels_btn = QPushButton("Export Labels")
|
||||
self.export_labels_btn.setToolTip(
|
||||
"Export YOLO .txt labels for the selected image/model run.\n"
|
||||
"Output path is inferred from the image path (images/ -> labels/)."
|
||||
)
|
||||
self.export_labels_btn.clicked.connect(self._export_labels_for_current_selection)
|
||||
controls_layout.addWidget(self.export_labels_btn)
|
||||
|
||||
controls_layout.addStretch()
|
||||
left_layout.addLayout(controls_layout)
|
||||
|
||||
self.results_table = QTableWidget(0, 5)
|
||||
self.results_table.setHorizontalHeaderLabels(
|
||||
["Image", "Model", "Detections", "Classes", "Last Updated"]
|
||||
)
|
||||
self.results_table.horizontalHeader().setSectionResizeMode(
|
||||
0, QHeaderView.Stretch
|
||||
)
|
||||
self.results_table.horizontalHeader().setSectionResizeMode(
|
||||
1, QHeaderView.Stretch
|
||||
)
|
||||
self.results_table.horizontalHeader().setSectionResizeMode(
|
||||
2, QHeaderView.ResizeToContents
|
||||
)
|
||||
self.results_table.horizontalHeader().setSectionResizeMode(
|
||||
3, QHeaderView.Stretch
|
||||
)
|
||||
self.results_table.horizontalHeader().setSectionResizeMode(
|
||||
4, QHeaderView.ResizeToContents
|
||||
)
|
||||
self.results_table.setHorizontalHeaderLabels(["Image", "Model", "Detections", "Classes", "Last Updated"])
|
||||
self.results_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.Stretch)
|
||||
self.results_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.Stretch)
|
||||
self.results_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeToContents)
|
||||
self.results_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.Stretch)
|
||||
self.results_table.horizontalHeader().setSectionResizeMode(4, QHeaderView.ResizeToContents)
|
||||
self.results_table.setSelectionBehavior(QAbstractItemView.SelectRows)
|
||||
self.results_table.setSelectionMode(QAbstractItemView.SingleSelection)
|
||||
self.results_table.setEditTriggers(QAbstractItemView.NoEditTriggers)
|
||||
@@ -106,6 +108,8 @@ class ResultsTab(QWidget):
|
||||
preview_layout = QVBoxLayout()
|
||||
|
||||
self.preview_canvas = AnnotationCanvasWidget()
|
||||
# Auto-zoom so newly loaded images fill the available preview viewport.
|
||||
self.preview_canvas.set_auto_fit_to_view(True)
|
||||
self.preview_canvas.set_polyline_enabled(False)
|
||||
self.preview_canvas.set_show_bboxes(True)
|
||||
preview_layout.addWidget(self.preview_canvas)
|
||||
@@ -119,9 +123,7 @@ class ResultsTab(QWidget):
|
||||
self.show_bboxes_checkbox.stateChanged.connect(self._toggle_bboxes)
|
||||
self.show_confidence_checkbox = QCheckBox("Show Confidence")
|
||||
self.show_confidence_checkbox.setChecked(False)
|
||||
self.show_confidence_checkbox.stateChanged.connect(
|
||||
self._apply_detection_overlays
|
||||
)
|
||||
self.show_confidence_checkbox.stateChanged.connect(self._apply_detection_overlays)
|
||||
toggles_layout.addWidget(self.show_masks_checkbox)
|
||||
toggles_layout.addWidget(self.show_bboxes_checkbox)
|
||||
toggles_layout.addWidget(self.show_confidence_checkbox)
|
||||
@@ -144,6 +146,41 @@ class ResultsTab(QWidget):
|
||||
layout.addWidget(splitter)
|
||||
self.setLayout(layout)
|
||||
|
||||
def _delete_all_detections(self):
|
||||
"""Delete all detections from the database after user confirmation."""
|
||||
confirm = QMessageBox.warning(
|
||||
self,
|
||||
"Delete All Detections",
|
||||
"This will permanently delete ALL detections from the database.\n\n"
|
||||
"This action cannot be undone.\n\n"
|
||||
"Do you want to continue?",
|
||||
QMessageBox.Yes | QMessageBox.No,
|
||||
QMessageBox.No,
|
||||
)
|
||||
|
||||
if confirm != QMessageBox.Yes:
|
||||
return
|
||||
|
||||
try:
|
||||
deleted = self.db_manager.delete_all_detections()
|
||||
except Exception as exc:
|
||||
logger.error(f"Failed to delete all detections: {exc}")
|
||||
QMessageBox.critical(
|
||||
self,
|
||||
"Error",
|
||||
f"Failed to delete detections:\n{exc}",
|
||||
)
|
||||
return
|
||||
|
||||
QMessageBox.information(
|
||||
self,
|
||||
"Delete All Detections",
|
||||
f"Deleted {deleted} detection(s) from the database.",
|
||||
)
|
||||
|
||||
# Reset UI state.
|
||||
self.refresh()
|
||||
|
||||
def refresh(self):
|
||||
"""Refresh the detection list and preview."""
|
||||
self._load_detection_summary()
|
||||
@@ -153,6 +190,8 @@ class ResultsTab(QWidget):
|
||||
self.current_detections = []
|
||||
self.preview_canvas.clear()
|
||||
self.summary_label.setText("Select a detection result to preview.")
|
||||
if hasattr(self, "export_labels_btn"):
|
||||
self.export_labels_btn.setEnabled(False)
|
||||
|
||||
def _load_detection_summary(self):
|
||||
"""Load latest detection summaries grouped by image + model."""
|
||||
@@ -169,8 +208,7 @@ class ResultsTab(QWidget):
|
||||
"image_id": det["image_id"],
|
||||
"model_id": det["model_id"],
|
||||
"image_path": det.get("image_path"),
|
||||
"image_filename": det.get("image_filename")
|
||||
or det.get("image_path"),
|
||||
"image_filename": det.get("image_filename") or det.get("image_path"),
|
||||
"model_name": det.get("model_name", ""),
|
||||
"model_version": det.get("model_version", ""),
|
||||
"last_detected": det.get("detected_at"),
|
||||
@@ -183,8 +221,7 @@ class ResultsTab(QWidget):
|
||||
|
||||
entry["count"] += 1
|
||||
if det.get("detected_at") and (
|
||||
not entry.get("last_detected")
|
||||
or str(det.get("detected_at")) > str(entry.get("last_detected"))
|
||||
not entry.get("last_detected") or str(det.get("detected_at")) > str(entry.get("last_detected"))
|
||||
):
|
||||
entry["last_detected"] = det.get("detected_at")
|
||||
if det.get("class_name"):
|
||||
@@ -214,9 +251,7 @@ class ResultsTab(QWidget):
|
||||
|
||||
for row, entry in enumerate(self.detection_summary):
|
||||
model_label = f"{entry['model_name']} {entry['model_version']}".strip()
|
||||
class_list = (
|
||||
", ".join(sorted(entry["classes"])) if entry["classes"] else "-"
|
||||
)
|
||||
class_list = ", ".join(sorted(entry["classes"])) if entry["classes"] else "-"
|
||||
|
||||
items = [
|
||||
QTableWidgetItem(entry.get("image_filename", "")),
|
||||
@@ -276,6 +311,231 @@ class ResultsTab(QWidget):
|
||||
self._load_detections_for_selection(entry)
|
||||
self._apply_detection_overlays()
|
||||
self._update_summary_label(entry)
|
||||
if hasattr(self, "export_labels_btn"):
|
||||
self.export_labels_btn.setEnabled(True)
|
||||
|
||||
def _export_labels_for_current_selection(self):
|
||||
"""Export YOLO label file(s) for the currently selected image/model."""
|
||||
if not self.current_selection:
|
||||
QMessageBox.information(self, "Export Labels", "Select a detection result first.")
|
||||
return
|
||||
|
||||
entry = self.current_selection
|
||||
|
||||
image_path_str = self._resolve_image_path(entry)
|
||||
if not image_path_str:
|
||||
QMessageBox.warning(
|
||||
self,
|
||||
"Export Labels",
|
||||
"Unable to locate the image file for this detection; cannot infer labels path.",
|
||||
)
|
||||
return
|
||||
|
||||
# Ensure we have the detections for the selection.
|
||||
if not self.current_detections:
|
||||
self._load_detections_for_selection(entry)
|
||||
|
||||
if not self.current_detections:
|
||||
QMessageBox.information(
|
||||
self,
|
||||
"Export Labels",
|
||||
"No detections found for this image/model selection.",
|
||||
)
|
||||
return
|
||||
|
||||
image_path = Path(image_path_str)
|
||||
try:
|
||||
label_path = self._infer_yolo_label_path(image_path)
|
||||
except Exception as exc:
|
||||
logger.error(f"Failed to infer label path for {image_path}: {exc}")
|
||||
QMessageBox.critical(
|
||||
self,
|
||||
"Export Labels",
|
||||
f"Failed to infer export path for labels:\n{exc}",
|
||||
)
|
||||
return
|
||||
|
||||
class_map = self._build_detection_class_index_map(self.current_detections)
|
||||
if not class_map:
|
||||
QMessageBox.warning(
|
||||
self,
|
||||
"Export Labels",
|
||||
"Unable to build class->index mapping (missing class names).",
|
||||
)
|
||||
return
|
||||
|
||||
lines_written = 0
|
||||
skipped = 0
|
||||
label_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
with open(label_path, "w", encoding="utf-8") as handle:
|
||||
print("writing to", label_path)
|
||||
for det in self.current_detections:
|
||||
yolo_line = self._format_detection_as_yolo_line(det, class_map)
|
||||
if not yolo_line:
|
||||
skipped += 1
|
||||
continue
|
||||
handle.write(yolo_line + "\n")
|
||||
lines_written += 1
|
||||
except OSError as exc:
|
||||
logger.error(f"Failed to write labels file {label_path}: {exc}")
|
||||
QMessageBox.critical(
|
||||
self,
|
||||
"Export Labels",
|
||||
f"Failed to write label file:\n{label_path}\n\n{exc}",
|
||||
)
|
||||
return
|
||||
|
||||
return
|
||||
# Optional: write a classes.txt next to the labels root to make the mapping discoverable.
|
||||
# This is not required by Ultralytics (data.yaml usually holds class names), but helps reuse.
|
||||
try:
|
||||
classes_txt = label_path.parent.parent / "classes.txt"
|
||||
classes_txt.parent.mkdir(parents=True, exist_ok=True)
|
||||
inv = {idx: name for name, idx in class_map.items()}
|
||||
with open(classes_txt, "w", encoding="utf-8") as handle:
|
||||
for idx in range(len(inv)):
|
||||
handle.write(f"{inv[idx]}\n")
|
||||
except Exception:
|
||||
# Non-fatal
|
||||
pass
|
||||
|
||||
QMessageBox.information(
|
||||
self,
|
||||
"Export Labels",
|
||||
f"Exported {lines_written} label line(s) to:\n{label_path}\n\nSkipped {skipped} invalid detection(s).",
|
||||
)
|
||||
|
||||
def _infer_yolo_label_path(self, image_path: Path) -> Path:
|
||||
"""Infer a YOLO label path from an image path.
|
||||
|
||||
If the image lives under an `images/` directory (anywhere in the path), we mirror the
|
||||
subpath under a sibling `labels/` directory at the same level.
|
||||
|
||||
Example:
|
||||
/dataset/train/images/sub/img.jpg -> /dataset/train/labels/sub/img.txt
|
||||
"""
|
||||
|
||||
resolved = image_path.expanduser().resolve()
|
||||
|
||||
# Find the nearest ancestor directory named 'images'
|
||||
images_dir: Optional[Path] = None
|
||||
for parent in [resolved.parent, *resolved.parents]:
|
||||
if parent.name.lower() == "images":
|
||||
images_dir = parent
|
||||
break
|
||||
|
||||
if images_dir is not None:
|
||||
rel = resolved.relative_to(images_dir)
|
||||
labels_dir = images_dir.parent / "labels"
|
||||
return (labels_dir / rel).with_suffix(".txt")
|
||||
|
||||
# Fallback: create a local sibling labels folder next to the image.
|
||||
return (resolved.parent / "labels" / resolved.name).with_suffix(".txt")
|
||||
|
||||
def _build_detection_class_index_map(self, detections: List[Dict]) -> Dict[str, int]:
|
||||
"""Build a stable class_name -> YOLO class index mapping.
|
||||
|
||||
Preference order:
|
||||
1) Database object_classes table (alphabetical class_name order)
|
||||
2) Fallback to class_name values present in the detections (alphabetical)
|
||||
"""
|
||||
|
||||
names: List[str] = []
|
||||
try:
|
||||
db_classes = self.db_manager.get_object_classes() or []
|
||||
names = [str(row.get("class_name")) for row in db_classes if row.get("class_name")]
|
||||
except Exception:
|
||||
names = []
|
||||
|
||||
if not names:
|
||||
observed = sorted({str(det.get("class_name")) for det in detections if det.get("class_name")})
|
||||
names = list(observed)
|
||||
|
||||
return {name: idx for idx, name in enumerate(names)}
|
||||
|
||||
def _format_detection_as_yolo_line(self, det: Dict, class_map: Dict[str, int]) -> Optional[str]:
|
||||
"""Convert a detection row to a YOLO label line.
|
||||
|
||||
- If segmentation_mask is present, exports segmentation polygon format:
|
||||
class x1 y1 x2 y2 ...
|
||||
(normalized coordinates)
|
||||
- Otherwise exports bbox format:
|
||||
class x_center y_center width height
|
||||
(normalized coordinates)
|
||||
"""
|
||||
|
||||
class_name = det.get("class_name")
|
||||
if not class_name or class_name not in class_map:
|
||||
return None
|
||||
class_idx = class_map[class_name]
|
||||
|
||||
mask = det.get("segmentation_mask")
|
||||
polygon = self._convert_segmentation_mask_to_polygon(mask)
|
||||
if polygon:
|
||||
coords = " ".join(f"{value:.6f}" for value in polygon)
|
||||
return f"{class_idx} {coords}".strip()
|
||||
|
||||
bbox = self._convert_bbox_to_yolo_xywh(det)
|
||||
if bbox is None:
|
||||
return None
|
||||
x_center, y_center, width, height = bbox
|
||||
return f"{class_idx} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
|
||||
|
||||
def _convert_bbox_to_yolo_xywh(self, det: Dict) -> Optional[Tuple[float, float, float, float]]:
|
||||
"""Convert stored xyxy (normalized) bbox to YOLO xywh (normalized)."""
|
||||
|
||||
x_min = det.get("x_min")
|
||||
y_min = det.get("y_min")
|
||||
x_max = det.get("x_max")
|
||||
y_max = det.get("y_max")
|
||||
if any(v is None for v in (x_min, y_min, x_max, y_max)):
|
||||
return None
|
||||
|
||||
try:
|
||||
x_min_f = self._clamp01(float(x_min))
|
||||
y_min_f = self._clamp01(float(y_min))
|
||||
x_max_f = self._clamp01(float(x_max))
|
||||
y_max_f = self._clamp01(float(y_max))
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
width = max(0.0, x_max_f - x_min_f)
|
||||
height = max(0.0, y_max_f - y_min_f)
|
||||
if width <= 0.0 or height <= 0.0:
|
||||
return None
|
||||
|
||||
x_center = x_min_f + width / 2.0
|
||||
y_center = y_min_f + height / 2.0
|
||||
return x_center, y_center, width, height
|
||||
|
||||
def _convert_segmentation_mask_to_polygon(self, mask_data) -> List[float]:
|
||||
"""Convert stored segmentation_mask [[x,y], ...] to YOLO polygon coords [x1,y1,...]."""
|
||||
|
||||
if not isinstance(mask_data, list):
|
||||
return []
|
||||
|
||||
coords: List[float] = []
|
||||
for point in mask_data:
|
||||
if not isinstance(point, (list, tuple)) or len(point) < 2:
|
||||
continue
|
||||
try:
|
||||
x = self._clamp01(float(point[0]))
|
||||
y = self._clamp01(float(point[1]))
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
coords.extend([x, y])
|
||||
|
||||
# Need at least 3 points => 6 values.
|
||||
return coords if len(coords) >= 6 else []
|
||||
|
||||
@staticmethod
|
||||
def _clamp01(value: float) -> float:
|
||||
if value < 0.0:
|
||||
return 0.0
|
||||
if value > 1.0:
|
||||
return 1.0
|
||||
return value
|
||||
|
||||
def _load_detections_for_selection(self, entry: Dict):
|
||||
"""Load detection records for the selected image/model pair."""
|
||||
|
||||
@@ -10,6 +10,7 @@ from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import yaml
|
||||
import numpy as np
|
||||
from PySide6.QtCore import Qt, QThread, Signal
|
||||
from PySide6.QtWidgets import (
|
||||
QWidget,
|
||||
@@ -34,7 +35,7 @@ from PySide6.QtWidgets import (
|
||||
from src.database.db_manager import DatabaseManager
|
||||
from src.model.yolo_wrapper import YOLOWrapper
|
||||
from src.utils.config_manager import ConfigManager
|
||||
from src.utils.image import Image, convert_grayscale_to_rgb_preserve_range
|
||||
from src.utils.image import Image
|
||||
from src.utils.logger import get_logger
|
||||
|
||||
|
||||
@@ -91,10 +92,7 @@ class TrainingWorker(QThread):
|
||||
},
|
||||
}
|
||||
]
|
||||
computed_total = sum(
|
||||
max(0, int((stage.get("params") or {}).get("epochs", 0)))
|
||||
for stage in self.stage_plan
|
||||
)
|
||||
computed_total = sum(max(0, int((stage.get("params") or {}).get("epochs", 0))) for stage in self.stage_plan)
|
||||
self.total_epochs = total_epochs if total_epochs else computed_total or epochs
|
||||
self._stop_requested = False
|
||||
|
||||
@@ -201,9 +199,7 @@ class TrainingWorker(QThread):
|
||||
class TrainingTab(QWidget):
|
||||
"""Training tab for model training."""
|
||||
|
||||
def __init__(
|
||||
self, db_manager: DatabaseManager, config_manager: ConfigManager, parent=None
|
||||
):
|
||||
def __init__(self, db_manager: DatabaseManager, config_manager: ConfigManager, parent=None):
|
||||
super().__init__(parent)
|
||||
self.db_manager = db_manager
|
||||
self.config_manager = config_manager
|
||||
@@ -337,18 +333,14 @@ class TrainingTab(QWidget):
|
||||
self.model_version_edit = QLineEdit("v1")
|
||||
form_layout.addRow("Version:", self.model_version_edit)
|
||||
|
||||
default_base_model = self.config_manager.get(
|
||||
"models.default_base_model", "yolov8s-seg.pt"
|
||||
)
|
||||
default_base_model = self.config_manager.get("models.default_base_model", "yolov8s-seg.pt")
|
||||
base_model_choices = self.config_manager.get("models.base_model_choices", [])
|
||||
|
||||
self.base_model_combo = QComboBox()
|
||||
self.base_model_combo.addItem("Custom path…", "")
|
||||
for choice in base_model_choices:
|
||||
self.base_model_combo.addItem(choice, choice)
|
||||
self.base_model_combo.currentIndexChanged.connect(
|
||||
self._on_base_model_preset_changed
|
||||
)
|
||||
self.base_model_combo.currentIndexChanged.connect(self._on_base_model_preset_changed)
|
||||
form_layout.addRow("Base Model Preset:", self.base_model_combo)
|
||||
|
||||
base_model_layout = QHBoxLayout()
|
||||
@@ -434,12 +426,8 @@ class TrainingTab(QWidget):
|
||||
group_layout = QVBoxLayout()
|
||||
|
||||
self.two_stage_checkbox = QCheckBox("Enable staged head-only + full fine-tune")
|
||||
two_stage_defaults = (
|
||||
training_defaults.get("two_stage", {}) if training_defaults else {}
|
||||
)
|
||||
self.two_stage_checkbox.setChecked(
|
||||
bool(two_stage_defaults.get("enabled", False))
|
||||
)
|
||||
two_stage_defaults = training_defaults.get("two_stage", {}) if training_defaults else {}
|
||||
self.two_stage_checkbox.setChecked(bool(two_stage_defaults.get("enabled", False)))
|
||||
self.two_stage_checkbox.toggled.connect(self._on_two_stage_toggled)
|
||||
group_layout.addWidget(self.two_stage_checkbox)
|
||||
|
||||
@@ -501,9 +489,7 @@ class TrainingTab(QWidget):
|
||||
stage2_group.setLayout(stage2_form)
|
||||
controls_layout.addWidget(stage2_group)
|
||||
|
||||
helper_label = QLabel(
|
||||
"When enabled, staged hyperparameters override the global epochs/patience/lr."
|
||||
)
|
||||
helper_label = QLabel("When enabled, staged hyperparameters override the global epochs/patience/lr.")
|
||||
helper_label.setWordWrap(True)
|
||||
controls_layout.addWidget(helper_label)
|
||||
|
||||
@@ -548,9 +534,7 @@ class TrainingTab(QWidget):
|
||||
if normalized == preset_value:
|
||||
target_index = idx
|
||||
break
|
||||
if normalized.endswith(f"/{preset_value}") or normalized.endswith(
|
||||
f"\\{preset_value}"
|
||||
):
|
||||
if normalized.endswith(f"/{preset_value}") or normalized.endswith(f"\\{preset_value}"):
|
||||
target_index = idx
|
||||
break
|
||||
self.base_model_combo.blockSignals(True)
|
||||
@@ -638,9 +622,7 @@ class TrainingTab(QWidget):
|
||||
|
||||
def _browse_dataset(self):
|
||||
"""Open a file dialog to manually select data.yaml."""
|
||||
start_dir = self.config_manager.get(
|
||||
"training.last_dataset_dir", "data/datasets"
|
||||
)
|
||||
start_dir = self.config_manager.get("training.last_dataset_dir", "data/datasets")
|
||||
start_path = Path(start_dir).expanduser()
|
||||
if not start_path.exists():
|
||||
start_path = Path.cwd()
|
||||
@@ -676,9 +658,7 @@ class TrainingTab(QWidget):
|
||||
return
|
||||
except Exception as exc:
|
||||
logger.exception("Unexpected error while generating data.yaml")
|
||||
self._display_dataset_error(
|
||||
"Unexpected error while generating data.yaml. Check logs for details."
|
||||
)
|
||||
self._display_dataset_error("Unexpected error while generating data.yaml. Check logs for details.")
|
||||
QMessageBox.critical(
|
||||
self,
|
||||
"data.yaml Generation Failed",
|
||||
@@ -755,13 +735,9 @@ class TrainingTab(QWidget):
|
||||
self.selected_dataset = info
|
||||
|
||||
self.dataset_root_label.setText(info["root"]) # type: ignore[arg-type]
|
||||
self.train_count_label.setText(
|
||||
self._format_split_info(info["splits"].get("train"))
|
||||
)
|
||||
self.train_count_label.setText(self._format_split_info(info["splits"].get("train")))
|
||||
self.val_count_label.setText(self._format_split_info(info["splits"].get("val")))
|
||||
self.test_count_label.setText(
|
||||
self._format_split_info(info["splits"].get("test"))
|
||||
)
|
||||
self.test_count_label.setText(self._format_split_info(info["splits"].get("test")))
|
||||
self.num_classes_label.setText(str(info["num_classes"]))
|
||||
class_names = ", ".join(info["class_names"]) or "–"
|
||||
self.class_names_label.setText(class_names)
|
||||
@@ -815,18 +791,12 @@ class TrainingTab(QWidget):
|
||||
if split_path.exists():
|
||||
split_info["count"] = self._count_images(split_path)
|
||||
if split_info["count"] == 0:
|
||||
warnings.append(
|
||||
f"No images found for {split_name} split at {split_path}"
|
||||
)
|
||||
warnings.append(f"No images found for {split_name} split at {split_path}")
|
||||
else:
|
||||
warnings.append(
|
||||
f"{split_name.capitalize()} path does not exist: {split_path}"
|
||||
)
|
||||
warnings.append(f"{split_name.capitalize()} path does not exist: {split_path}")
|
||||
else:
|
||||
if split_name in ("train", "val"):
|
||||
warnings.append(
|
||||
f"{split_name.capitalize()} split missing in data.yaml"
|
||||
)
|
||||
warnings.append(f"{split_name.capitalize()} split missing in data.yaml")
|
||||
splits[split_name] = split_info
|
||||
|
||||
names_list = self._normalize_class_names(data.get("names"))
|
||||
@@ -844,9 +814,7 @@ class TrainingTab(QWidget):
|
||||
if not names_list and nc_value:
|
||||
names_list = [f"class_{idx}" for idx in range(int(nc_value))]
|
||||
elif nc_value and len(names_list) not in (0, int(nc_value)):
|
||||
warnings.append(
|
||||
f"Number of class names ({len(names_list)}) does not match nc={nc_value}"
|
||||
)
|
||||
warnings.append(f"Number of class names ({len(names_list)}) does not match nc={nc_value}")
|
||||
|
||||
dataset_name = data.get("name") or base_path.name
|
||||
|
||||
@@ -898,16 +866,12 @@ class TrainingTab(QWidget):
|
||||
|
||||
class_index_map = self._build_class_index_map(dataset_info)
|
||||
if not class_index_map:
|
||||
self._append_training_log(
|
||||
"Skipping label export: dataset classes do not match database entries."
|
||||
)
|
||||
self._append_training_log("Skipping label export: dataset classes do not match database entries.")
|
||||
return
|
||||
|
||||
dataset_root_str = dataset_info.get("root")
|
||||
dataset_yaml_path = dataset_info.get("yaml_path")
|
||||
dataset_yaml = (
|
||||
Path(dataset_yaml_path).expanduser() if dataset_yaml_path else None
|
||||
)
|
||||
dataset_yaml = Path(dataset_yaml_path).expanduser() if dataset_yaml_path else None
|
||||
dataset_root: Optional[Path]
|
||||
if dataset_root_str:
|
||||
dataset_root = Path(dataset_root_str).resolve()
|
||||
@@ -941,7 +905,9 @@ class TrainingTab(QWidget):
|
||||
if stats["registered_images"]:
|
||||
message += f" {stats['registered_images']} image(s) had database-backed annotations."
|
||||
if stats["missing_records"]:
|
||||
message += f" {stats['missing_records']} image(s) had no database entry; empty label files were written."
|
||||
message += (
|
||||
f" {stats['missing_records']} image(s) had no database entry; empty label files were written."
|
||||
)
|
||||
split_messages.append(message)
|
||||
|
||||
for msg in split_messages:
|
||||
@@ -973,9 +939,7 @@ class TrainingTab(QWidget):
|
||||
continue
|
||||
|
||||
processed_images += 1
|
||||
label_path = (labels_dir / image_file.relative_to(images_dir)).with_suffix(
|
||||
".txt"
|
||||
)
|
||||
label_path = (labels_dir / image_file.relative_to(images_dir)).with_suffix(".txt")
|
||||
label_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
found, annotation_entries = self._fetch_annotations_for_image(
|
||||
@@ -991,25 +955,23 @@ class TrainingTab(QWidget):
|
||||
for entry in annotation_entries:
|
||||
polygon = entry.get("polygon") or []
|
||||
if polygon:
|
||||
print(image_file, polygon[:4], polygon[-2:], entry.get("bbox"))
|
||||
# coords = " ".join(f"{value:.6f}" for value in entry.get("bbox"))
|
||||
# coords += " "
|
||||
coords = " ".join(f"{value:.6f}" for value in polygon)
|
||||
handle.write(f"{entry['class_idx']} {coords}\n")
|
||||
annotations_written += 1
|
||||
elif entry.get("bbox"):
|
||||
x_center, y_center, width, height = entry["bbox"]
|
||||
handle.write(
|
||||
f"{entry['class_idx']} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n"
|
||||
)
|
||||
handle.write(f"{entry['class_idx']} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
|
||||
annotations_written += 1
|
||||
|
||||
total_annotations += annotations_written
|
||||
|
||||
cache_reset_root = labels_dir.parent
|
||||
self._invalidate_split_cache(cache_reset_root)
|
||||
|
||||
if processed_images == 0:
|
||||
self._append_training_log(
|
||||
f"[{split_name}] No images found to export labels for."
|
||||
)
|
||||
self._append_training_log(f"[{split_name}] No images found to export labels for.")
|
||||
return None
|
||||
|
||||
return {
|
||||
@@ -1135,6 +1097,10 @@ class TrainingTab(QWidget):
|
||||
xs.append(x_val)
|
||||
ys.append(y_val)
|
||||
|
||||
if any(np.abs(np.array(coords[:2]) - np.array(coords[-2:])) < 1e-5):
|
||||
print("Closing polygon")
|
||||
coords.extend(coords[:2])
|
||||
|
||||
if len(coords) < 6:
|
||||
continue
|
||||
|
||||
@@ -1147,6 +1113,11 @@ class TrainingTab(QWidget):
|
||||
+ abs((min(ys) if ys else 0.0) - y_min)
|
||||
+ abs((max(ys) if ys else 0.0) - y_max)
|
||||
)
|
||||
width = max(0.0, x_max - x_min)
|
||||
height = max(0.0, y_max - y_min)
|
||||
x_center = x_min + width / 2.0
|
||||
y_center = y_min + height / 2.0
|
||||
score = (x_center, y_center, width, height)
|
||||
|
||||
candidates.append((score, coords))
|
||||
|
||||
@@ -1164,13 +1135,10 @@ class TrainingTab(QWidget):
|
||||
return 1.0
|
||||
return value
|
||||
|
||||
def _prepare_dataset_for_training(
|
||||
self, dataset_yaml: Path, dataset_info: Optional[Dict[str, Any]] = None
|
||||
) -> Path:
|
||||
def _prepare_dataset_for_training(self, dataset_yaml: Path, dataset_info: Optional[Dict[str, Any]] = None) -> Path:
|
||||
dataset_info = dataset_info or (
|
||||
self.selected_dataset
|
||||
if self.selected_dataset
|
||||
and self.selected_dataset.get("yaml_path") == str(dataset_yaml)
|
||||
if self.selected_dataset and self.selected_dataset.get("yaml_path") == str(dataset_yaml)
|
||||
else self._parse_dataset_yaml(dataset_yaml)
|
||||
)
|
||||
|
||||
@@ -1189,14 +1157,10 @@ class TrainingTab(QWidget):
|
||||
cache_root = self._get_rgb_cache_root(dataset_yaml)
|
||||
rgb_yaml = cache_root / "data.yaml"
|
||||
if rgb_yaml.exists():
|
||||
self._append_training_log(
|
||||
f"Detected grayscale dataset; reusing RGB cache at {cache_root}"
|
||||
)
|
||||
self._append_training_log(f"Detected grayscale dataset; reusing RGB cache at {cache_root}")
|
||||
return rgb_yaml
|
||||
|
||||
self._append_training_log(
|
||||
f"Detected grayscale dataset; creating RGB cache at {cache_root}"
|
||||
)
|
||||
self._append_training_log(f"Detected grayscale dataset; creating RGB cache at {cache_root}")
|
||||
self._build_rgb_dataset(cache_root, dataset_info)
|
||||
return rgb_yaml
|
||||
|
||||
@@ -1303,6 +1267,14 @@ class TrainingTab(QWidget):
|
||||
sample_image = self._find_first_image(images_dir)
|
||||
if not sample_image:
|
||||
return False
|
||||
|
||||
# Do not force an RGB cache for TIFF datasets.
|
||||
# We handle grayscale/16-bit TIFFs via runtime Ultralytics patches that:
|
||||
# - load TIFFs with `tifffile`
|
||||
# - replicate grayscale to 3 channels without quantization
|
||||
# - normalize uint16 correctly during training
|
||||
if sample_image.suffix.lower() in {".tif", ".tiff"}:
|
||||
return False
|
||||
try:
|
||||
img = Image(sample_image)
|
||||
return img.pil_image.mode.upper() != "RGB"
|
||||
@@ -1368,7 +1340,7 @@ class TrainingTab(QWidget):
|
||||
img_obj = Image(src)
|
||||
pil_img = img_obj.pil_image
|
||||
if len(pil_img.getbands()) == 1:
|
||||
rgb_img = convert_grayscale_to_rgb_preserve_range(pil_img)
|
||||
rgb_img = img_obj.convert_grayscale_to_rgb_preserve_range()
|
||||
else:
|
||||
rgb_img = pil_img.convert("RGB")
|
||||
rgb_img.save(dst)
|
||||
@@ -1455,15 +1427,12 @@ class TrainingTab(QWidget):
|
||||
|
||||
dataset_path = Path(dataset_yaml).expanduser()
|
||||
if not dataset_path.exists():
|
||||
QMessageBox.warning(
|
||||
self, "Invalid Dataset", "Selected data.yaml file does not exist."
|
||||
)
|
||||
QMessageBox.warning(self, "Invalid Dataset", "Selected data.yaml file does not exist.")
|
||||
return
|
||||
|
||||
dataset_info = (
|
||||
self.selected_dataset
|
||||
if self.selected_dataset
|
||||
and self.selected_dataset.get("yaml_path") == str(dataset_path)
|
||||
if self.selected_dataset and self.selected_dataset.get("yaml_path") == str(dataset_path)
|
||||
else self._parse_dataset_yaml(dataset_path)
|
||||
)
|
||||
|
||||
@@ -1472,16 +1441,12 @@ class TrainingTab(QWidget):
|
||||
|
||||
dataset_to_use = self._prepare_dataset_for_training(dataset_path, dataset_info)
|
||||
if dataset_to_use != dataset_path:
|
||||
self._append_training_log(
|
||||
f"Using RGB-converted dataset at {dataset_to_use.parent}"
|
||||
)
|
||||
self._append_training_log(f"Using RGB-converted dataset at {dataset_to_use.parent}")
|
||||
|
||||
params = self._collect_training_params()
|
||||
stage_plan = self._compose_stage_plan(params)
|
||||
params["stage_plan"] = stage_plan
|
||||
total_planned_epochs = (
|
||||
self._calculate_total_stage_epochs(stage_plan) or params["epochs"]
|
||||
)
|
||||
total_planned_epochs = self._calculate_total_stage_epochs(stage_plan) or params["epochs"]
|
||||
params["total_planned_epochs"] = total_planned_epochs
|
||||
self._active_training_params = params
|
||||
self._training_cancelled = False
|
||||
@@ -1490,9 +1455,7 @@ class TrainingTab(QWidget):
|
||||
self._append_training_log("Two-stage fine-tuning schedule:")
|
||||
self._log_stage_plan(stage_plan)
|
||||
|
||||
self._append_training_log(
|
||||
f"Starting training run '{params['run_name']}' using {params['base_model']}"
|
||||
)
|
||||
self._append_training_log(f"Starting training run '{params['run_name']}' using {params['base_model']}")
|
||||
|
||||
self.training_progress_bar.setVisible(True)
|
||||
self.training_progress_bar.setMaximum(max(1, total_planned_epochs))
|
||||
@@ -1520,9 +1483,7 @@ class TrainingTab(QWidget):
|
||||
def _stop_training(self):
|
||||
if self.training_worker and self.training_worker.isRunning():
|
||||
self._training_cancelled = True
|
||||
self._append_training_log(
|
||||
"Stop requested. Waiting for the current epoch to finish..."
|
||||
)
|
||||
self._append_training_log("Stop requested. Waiting for the current epoch to finish...")
|
||||
self.training_worker.stop()
|
||||
self.stop_training_button.setEnabled(False)
|
||||
|
||||
@@ -1558,9 +1519,7 @@ class TrainingTab(QWidget):
|
||||
|
||||
if worker.isRunning():
|
||||
if not worker.wait(wait_timeout_ms):
|
||||
logger.warning(
|
||||
"Training worker did not finish within %sms", wait_timeout_ms
|
||||
)
|
||||
logger.warning("Training worker did not finish within %sms", wait_timeout_ms)
|
||||
|
||||
worker.deleteLater()
|
||||
|
||||
@@ -1577,16 +1536,12 @@ class TrainingTab(QWidget):
|
||||
self._set_training_state(False)
|
||||
self.training_progress_bar.setVisible(False)
|
||||
|
||||
def _on_training_progress(
|
||||
self, current_epoch: int, total_epochs: int, metrics: Dict[str, Any]
|
||||
):
|
||||
def _on_training_progress(self, current_epoch: int, total_epochs: int, metrics: Dict[str, Any]):
|
||||
self.training_progress_bar.setMaximum(total_epochs)
|
||||
self.training_progress_bar.setValue(current_epoch)
|
||||
parts = [f"Epoch {current_epoch}/{total_epochs}"]
|
||||
if metrics:
|
||||
metric_text = ", ".join(
|
||||
f"{key}: {value:.4f}" for key, value in metrics.items()
|
||||
)
|
||||
metric_text = ", ".join(f"{key}: {value:.4f}" for key, value in metrics.items())
|
||||
parts.append(metric_text)
|
||||
self._append_training_log(" | ".join(parts))
|
||||
|
||||
@@ -1613,9 +1568,7 @@ class TrainingTab(QWidget):
|
||||
f"Model trained but not registered: {exc}",
|
||||
)
|
||||
else:
|
||||
QMessageBox.information(
|
||||
self, "Training Complete", "Training finished successfully."
|
||||
)
|
||||
QMessageBox.information(self, "Training Complete", "Training finished successfully.")
|
||||
|
||||
def _on_training_error(self, message: str):
|
||||
self._cleanup_training_worker()
|
||||
@@ -1661,9 +1614,7 @@ class TrainingTab(QWidget):
|
||||
metrics=results.get("metrics"),
|
||||
)
|
||||
|
||||
self._append_training_log(
|
||||
f"Registered model '{params['model_name']}' (ID {model_id}) at {model_path}"
|
||||
)
|
||||
self._append_training_log(f"Registered model '{params['model_name']}' (ID {model_id}) at {model_path}")
|
||||
self._active_training_params = None
|
||||
|
||||
def _set_training_state(self, is_training: bool):
|
||||
@@ -1706,9 +1657,7 @@ class TrainingTab(QWidget):
|
||||
|
||||
def _browse_save_dir(self):
|
||||
start_path = self.save_dir_edit.text().strip() or "data/models"
|
||||
directory = QFileDialog.getExistingDirectory(
|
||||
self, "Select Save Directory", start_path
|
||||
)
|
||||
directory = QFileDialog.getExistingDirectory(self, "Select Save Directory", start_path)
|
||||
if directory:
|
||||
self.save_dir_edit.setText(directory)
|
||||
|
||||
|
||||
@@ -2,45 +2,554 @@
|
||||
Validation tab for the microscopy object detection application.
|
||||
"""
|
||||
|
||||
from PySide6.QtWidgets import QWidget, QVBoxLayout, QLabel, QGroupBox
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
||||
|
||||
from PySide6.QtCore import Qt, QSize
|
||||
from PySide6.QtGui import QPainter, QPixmap
|
||||
from PySide6.QtWidgets import (
|
||||
QWidget,
|
||||
QVBoxLayout,
|
||||
QLabel,
|
||||
QGroupBox,
|
||||
QHBoxLayout,
|
||||
QPushButton,
|
||||
QComboBox,
|
||||
QFormLayout,
|
||||
QScrollArea,
|
||||
QGridLayout,
|
||||
QFrame,
|
||||
QTableWidget,
|
||||
QTableWidgetItem,
|
||||
QHeaderView,
|
||||
QSplitter,
|
||||
QListWidget,
|
||||
QListWidgetItem,
|
||||
QAbstractItemView,
|
||||
QGraphicsView,
|
||||
QGraphicsScene,
|
||||
QGraphicsPixmapItem,
|
||||
)
|
||||
|
||||
from src.database.db_manager import DatabaseManager
|
||||
from src.utils.config_manager import ConfigManager
|
||||
from src.utils.logger import get_logger
|
||||
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _PlotItem:
|
||||
label: str
|
||||
path: Path
|
||||
|
||||
|
||||
class _ZoomableImageView(QGraphicsView):
|
||||
"""Zoomable image viewer.
|
||||
|
||||
- Mouse wheel: zoom in/out
|
||||
- Left mouse drag: pan (ScrollHandDrag)
|
||||
"""
|
||||
|
||||
def __init__(self, parent: Optional[QWidget] = None):
|
||||
super().__init__(parent)
|
||||
self._scene = QGraphicsScene(self)
|
||||
self.setScene(self._scene)
|
||||
self._pixmap_item = QGraphicsPixmapItem()
|
||||
self._scene.addItem(self._pixmap_item)
|
||||
|
||||
# QGraphicsView render hints are QPainter.RenderHints.
|
||||
self.setRenderHints(self.renderHints() | QPainter.RenderHint.SmoothPixmapTransform)
|
||||
self.setDragMode(QGraphicsView.DragMode.ScrollHandDrag)
|
||||
self.setTransformationAnchor(QGraphicsView.ViewportAnchor.AnchorUnderMouse)
|
||||
self.setResizeAnchor(QGraphicsView.ViewportAnchor.AnchorUnderMouse)
|
||||
|
||||
self._has_pixmap = False
|
||||
|
||||
def clear(self) -> None:
|
||||
self._pixmap_item.setPixmap(QPixmap())
|
||||
self._scene.setSceneRect(0, 0, 1, 1)
|
||||
self.resetTransform()
|
||||
self._has_pixmap = False
|
||||
|
||||
def set_pixmap(self, pixmap: QPixmap, *, fit: bool = True) -> None:
|
||||
self._pixmap_item.setPixmap(pixmap)
|
||||
self._scene.setSceneRect(pixmap.rect())
|
||||
self._has_pixmap = not pixmap.isNull()
|
||||
self.resetTransform()
|
||||
if fit and self._has_pixmap:
|
||||
self.fitInView(self._pixmap_item, Qt.AspectRatioMode.KeepAspectRatio)
|
||||
|
||||
def wheelEvent(self, event) -> None: # type: ignore[override]
|
||||
if not self._has_pixmap:
|
||||
return
|
||||
zoom_in_factor = 1.25
|
||||
zoom_out_factor = 1.0 / zoom_in_factor
|
||||
factor = zoom_in_factor if event.angleDelta().y() > 0 else zoom_out_factor
|
||||
self.scale(factor, factor)
|
||||
|
||||
|
||||
class ValidationTab(QWidget):
|
||||
"""Validation tab placeholder."""
|
||||
"""Validation tab that shows stored validation metrics + plots for a selected model."""
|
||||
|
||||
def __init__(
|
||||
self, db_manager: DatabaseManager, config_manager: ConfigManager, parent=None
|
||||
):
|
||||
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._models: List[Dict[str, Any]] = []
|
||||
self._selected_model_id: Optional[int] = None
|
||||
self._plot_widgets: List[QWidget] = []
|
||||
self._plot_items: List[_PlotItem] = []
|
||||
|
||||
self._setup_ui()
|
||||
self.refresh()
|
||||
|
||||
def _setup_ui(self):
|
||||
"""Setup user interface."""
|
||||
layout = QVBoxLayout()
|
||||
layout = QVBoxLayout(self)
|
||||
|
||||
group = QGroupBox("Validation")
|
||||
group_layout = QVBoxLayout()
|
||||
label = QLabel(
|
||||
"Validation functionality will be implemented here.\n\n"
|
||||
"Features:\n"
|
||||
"- Model validation\n"
|
||||
"- Metrics visualization\n"
|
||||
"- Confusion matrix\n"
|
||||
"- Precision-Recall curves"
|
||||
)
|
||||
group_layout.addWidget(label)
|
||||
group.setLayout(group_layout)
|
||||
# ===== Header controls =====
|
||||
header = QGroupBox("Validation")
|
||||
header_layout = QVBoxLayout()
|
||||
header_row = QHBoxLayout()
|
||||
|
||||
layout.addWidget(group)
|
||||
layout.addStretch()
|
||||
self.setLayout(layout)
|
||||
header_row.addWidget(QLabel("Select model:"))
|
||||
|
||||
self.model_combo = QComboBox()
|
||||
self.model_combo.setMinimumWidth(420)
|
||||
self.model_combo.currentIndexChanged.connect(self._on_model_selected)
|
||||
header_row.addWidget(self.model_combo, 1)
|
||||
|
||||
self.refresh_btn = QPushButton("Refresh")
|
||||
self.refresh_btn.clicked.connect(self.refresh)
|
||||
header_row.addWidget(self.refresh_btn)
|
||||
header_row.addStretch()
|
||||
|
||||
header_layout.addLayout(header_row)
|
||||
self.header_status = QLabel("No models loaded.")
|
||||
self.header_status.setWordWrap(True)
|
||||
header_layout.addWidget(self.header_status)
|
||||
header.setLayout(header_layout)
|
||||
layout.addWidget(header)
|
||||
|
||||
# ===== Metrics =====
|
||||
metrics_group = QGroupBox("Validation Metrics")
|
||||
metrics_layout = QVBoxLayout()
|
||||
|
||||
self.metrics_form = QFormLayout()
|
||||
self.metric_labels: Dict[str, QLabel] = {}
|
||||
for key in ("mAP50", "mAP50-95", "precision", "recall", "fitness"):
|
||||
value_label = QLabel("–")
|
||||
value_label.setTextInteractionFlags(Qt.TextSelectableByMouse)
|
||||
self.metric_labels[key] = value_label
|
||||
self.metrics_form.addRow(f"{key}:", value_label)
|
||||
metrics_layout.addLayout(self.metrics_form)
|
||||
|
||||
self.per_class_table = QTableWidget(0, 3)
|
||||
self.per_class_table.setHorizontalHeaderLabels(["Class", "AP", "AP50"])
|
||||
self.per_class_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.Stretch)
|
||||
self.per_class_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeToContents)
|
||||
self.per_class_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeToContents)
|
||||
self.per_class_table.setEditTriggers(QTableWidget.NoEditTriggers)
|
||||
self.per_class_table.setMinimumHeight(160)
|
||||
metrics_layout.addWidget(QLabel("Per-class metrics (if available):"))
|
||||
metrics_layout.addWidget(self.per_class_table)
|
||||
|
||||
metrics_group.setLayout(metrics_layout)
|
||||
layout.addWidget(metrics_group)
|
||||
|
||||
# ===== Plots =====
|
||||
plots_group = QGroupBox("Validation Plots")
|
||||
plots_layout = QVBoxLayout()
|
||||
|
||||
self.plots_status = QLabel("Select a model to see validation plots.")
|
||||
self.plots_status.setWordWrap(True)
|
||||
plots_layout.addWidget(self.plots_status)
|
||||
|
||||
self.plots_splitter = QSplitter(Qt.Orientation.Horizontal)
|
||||
|
||||
# Left: selected image viewer
|
||||
left_widget = QWidget()
|
||||
left_layout = QVBoxLayout(left_widget)
|
||||
left_layout.setContentsMargins(0, 0, 0, 0)
|
||||
|
||||
self.selected_plot_title = QLabel("No image selected.")
|
||||
self.selected_plot_title.setWordWrap(True)
|
||||
self.selected_plot_title.setTextInteractionFlags(Qt.TextSelectableByMouse)
|
||||
left_layout.addWidget(self.selected_plot_title)
|
||||
|
||||
self.plot_view = _ZoomableImageView()
|
||||
self.plot_view.setMinimumHeight(360)
|
||||
left_layout.addWidget(self.plot_view, 1)
|
||||
|
||||
self.selected_plot_path = QLabel("")
|
||||
self.selected_plot_path.setWordWrap(True)
|
||||
self.selected_plot_path.setStyleSheet("color: #888;")
|
||||
self.selected_plot_path.setTextInteractionFlags(Qt.TextSelectableByMouse)
|
||||
left_layout.addWidget(self.selected_plot_path)
|
||||
|
||||
# Right: scrollable list
|
||||
right_widget = QWidget()
|
||||
right_layout = QVBoxLayout(right_widget)
|
||||
right_layout.setContentsMargins(0, 0, 0, 0)
|
||||
right_layout.addWidget(QLabel("Images:"))
|
||||
|
||||
self.plots_list = QListWidget()
|
||||
self.plots_list.setSelectionMode(QAbstractItemView.SelectionMode.SingleSelection)
|
||||
self.plots_list.setIconSize(QSize(160, 160))
|
||||
self.plots_list.itemSelectionChanged.connect(self._on_plot_item_selected)
|
||||
right_layout.addWidget(self.plots_list, 1)
|
||||
|
||||
self.plots_splitter.addWidget(left_widget)
|
||||
self.plots_splitter.addWidget(right_widget)
|
||||
self.plots_splitter.setStretchFactor(0, 3)
|
||||
self.plots_splitter.setStretchFactor(1, 1)
|
||||
plots_layout.addWidget(self.plots_splitter, 1)
|
||||
|
||||
plots_group.setLayout(plots_layout)
|
||||
layout.addWidget(plots_group, 1)
|
||||
|
||||
layout.addStretch(0)
|
||||
|
||||
self._clear_metrics()
|
||||
self._clear_plots()
|
||||
|
||||
# ==================== Public API ====================
|
||||
|
||||
def refresh(self):
|
||||
"""Refresh the tab."""
|
||||
self._load_models()
|
||||
self._populate_model_combo()
|
||||
self._restore_or_select_default_model()
|
||||
|
||||
# ==================== Internal: models ====================
|
||||
|
||||
def _load_models(self) -> None:
|
||||
try:
|
||||
self._models = self.db_manager.get_models() or []
|
||||
except Exception as exc:
|
||||
logger.error("Failed to load models: %s", exc)
|
||||
self._models = []
|
||||
|
||||
def _populate_model_combo(self) -> None:
|
||||
self.model_combo.blockSignals(True)
|
||||
self.model_combo.clear()
|
||||
self.model_combo.addItem("Select a model…", None)
|
||||
|
||||
for model in self._models:
|
||||
model_id = model.get("id")
|
||||
name = (model.get("model_name") or "").strip()
|
||||
version = (model.get("model_version") or "").strip()
|
||||
created_at = model.get("created_at")
|
||||
label = f"{name} {version}".strip()
|
||||
if created_at:
|
||||
label = f"{label} ({created_at})"
|
||||
self.model_combo.addItem(label, model_id)
|
||||
|
||||
self.model_combo.blockSignals(False)
|
||||
|
||||
if self._models:
|
||||
self.header_status.setText(f"Loaded {len(self._models)} model(s).")
|
||||
else:
|
||||
self.header_status.setText("No models found. Train a model first.")
|
||||
|
||||
def _restore_or_select_default_model(self) -> None:
|
||||
if not self._models:
|
||||
self._selected_model_id = None
|
||||
self._clear_metrics()
|
||||
self._clear_plots()
|
||||
return
|
||||
|
||||
# Keep selection if still present.
|
||||
if self._selected_model_id is not None:
|
||||
for idx in range(1, self.model_combo.count()):
|
||||
if self.model_combo.itemData(idx) == self._selected_model_id:
|
||||
self.model_combo.setCurrentIndex(idx)
|
||||
return
|
||||
|
||||
# Otherwise select the newest model (top of get_models ORDER BY created_at DESC).
|
||||
first_model_id = self.model_combo.itemData(1) if self.model_combo.count() > 1 else None
|
||||
if first_model_id is not None:
|
||||
self.model_combo.setCurrentIndex(1)
|
||||
|
||||
def _on_model_selected(self, index: int) -> None:
|
||||
model_id = self.model_combo.itemData(index)
|
||||
if not model_id:
|
||||
self._selected_model_id = None
|
||||
self._clear_metrics()
|
||||
self._clear_plots()
|
||||
self.plots_status.setText("Select a model to see validation plots.")
|
||||
return
|
||||
|
||||
self._selected_model_id = int(model_id)
|
||||
model = self._get_model_by_id(self._selected_model_id)
|
||||
if not model:
|
||||
self._clear_metrics()
|
||||
self._clear_plots()
|
||||
self.plots_status.setText("Selected model not found.")
|
||||
return
|
||||
|
||||
self._render_metrics(model)
|
||||
self._render_plots(model)
|
||||
|
||||
def _get_model_by_id(self, model_id: int) -> Optional[Dict[str, Any]]:
|
||||
for model in self._models:
|
||||
if model.get("id") == model_id:
|
||||
return model
|
||||
try:
|
||||
return self.db_manager.get_model_by_id(model_id)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
# ==================== Internal: metrics ====================
|
||||
|
||||
def _clear_metrics(self) -> None:
|
||||
for label in self.metric_labels.values():
|
||||
label.setText("–")
|
||||
self.per_class_table.setRowCount(0)
|
||||
|
||||
def _render_metrics(self, model: Dict[str, Any]) -> None:
|
||||
self._clear_metrics()
|
||||
|
||||
metrics: Dict[str, Any] = model.get("metrics") or {}
|
||||
# Training tab stores metrics under results['metrics'] in training results payload.
|
||||
if isinstance(metrics, dict) and "metrics" in metrics and isinstance(metrics.get("metrics"), dict):
|
||||
metrics = metrics.get("metrics") or {}
|
||||
|
||||
def set_metric(key: str, value: Any) -> None:
|
||||
if key not in self.metric_labels:
|
||||
return
|
||||
if value is None:
|
||||
self.metric_labels[key].setText("–")
|
||||
return
|
||||
try:
|
||||
self.metric_labels[key].setText(f"{float(value):.4f}")
|
||||
except Exception:
|
||||
self.metric_labels[key].setText(str(value))
|
||||
|
||||
set_metric("mAP50", metrics.get("mAP50"))
|
||||
set_metric("mAP50-95", metrics.get("mAP50-95") or metrics.get("mAP50_95") or metrics.get("mAP50-95"))
|
||||
set_metric("precision", metrics.get("precision"))
|
||||
set_metric("recall", metrics.get("recall"))
|
||||
set_metric("fitness", metrics.get("fitness"))
|
||||
|
||||
# Optional per-class metrics
|
||||
class_metrics = metrics.get("class_metrics") if isinstance(metrics, dict) else None
|
||||
if isinstance(class_metrics, dict) and class_metrics:
|
||||
items = sorted(class_metrics.items(), key=lambda kv: str(kv[0]))
|
||||
self.per_class_table.setRowCount(len(items))
|
||||
for row, (cls_name, cls_stats) in enumerate(items):
|
||||
ap = (cls_stats or {}).get("ap")
|
||||
ap50 = (cls_stats or {}).get("ap50")
|
||||
self.per_class_table.setItem(row, 0, QTableWidgetItem(str(cls_name)))
|
||||
self.per_class_table.setItem(row, 1, QTableWidgetItem(self._format_float(ap)))
|
||||
self.per_class_table.setItem(row, 2, QTableWidgetItem(self._format_float(ap50)))
|
||||
else:
|
||||
self.per_class_table.setRowCount(0)
|
||||
|
||||
@staticmethod
|
||||
def _format_float(value: Any) -> str:
|
||||
if value is None:
|
||||
return "–"
|
||||
try:
|
||||
return f"{float(value):.4f}"
|
||||
except Exception:
|
||||
return str(value)
|
||||
|
||||
# ==================== Internal: plots ====================
|
||||
|
||||
def _clear_plots(self) -> None:
|
||||
# Remove legacy grid widgets (from the initial implementation).
|
||||
for widget in self._plot_widgets:
|
||||
widget.setParent(None)
|
||||
widget.deleteLater()
|
||||
self._plot_widgets = []
|
||||
|
||||
self._plot_items = []
|
||||
|
||||
if hasattr(self, "plots_list"):
|
||||
self.plots_list.blockSignals(True)
|
||||
self.plots_list.clear()
|
||||
self.plots_list.blockSignals(False)
|
||||
|
||||
if hasattr(self, "plot_view"):
|
||||
self.plot_view.clear()
|
||||
if hasattr(self, "selected_plot_title"):
|
||||
self.selected_plot_title.setText("No image selected.")
|
||||
if hasattr(self, "selected_plot_path"):
|
||||
self.selected_plot_path.setText("")
|
||||
|
||||
def _render_plots(self, model: Dict[str, Any]) -> None:
|
||||
self._clear_plots()
|
||||
|
||||
plot_dirs = self._infer_run_directories(model)
|
||||
plot_items = self._discover_plot_items(plot_dirs)
|
||||
|
||||
if not plot_items:
|
||||
dirs_text = "\n".join(str(p) for p in plot_dirs if p)
|
||||
self.plots_status.setText(
|
||||
"No validation plot images found for this model.\n\n"
|
||||
"Searched directories:\n" + (dirs_text or "(none)")
|
||||
)
|
||||
return
|
||||
|
||||
self._plot_items = list(plot_items)
|
||||
self.plots_status.setText(f"Found {len(plot_items)} plot image(s). Select one to view/zoom.")
|
||||
|
||||
self.plots_list.blockSignals(True)
|
||||
self.plots_list.clear()
|
||||
for idx, item in enumerate(self._plot_items):
|
||||
qitem = QListWidgetItem(item.label)
|
||||
qitem.setData(Qt.ItemDataRole.UserRole, idx)
|
||||
|
||||
pix = QPixmap(str(item.path))
|
||||
if not pix.isNull():
|
||||
thumb = pix.scaled(
|
||||
self.plots_list.iconSize(),
|
||||
Qt.AspectRatioMode.KeepAspectRatio,
|
||||
Qt.TransformationMode.SmoothTransformation,
|
||||
)
|
||||
qitem.setIcon(thumb)
|
||||
self.plots_list.addItem(qitem)
|
||||
self.plots_list.blockSignals(False)
|
||||
|
||||
if self.plots_list.count() > 0:
|
||||
self.plots_list.setCurrentRow(0)
|
||||
|
||||
def _on_plot_item_selected(self) -> None:
|
||||
if not self._plot_items:
|
||||
return
|
||||
|
||||
selected = self.plots_list.selectedItems()
|
||||
if not selected:
|
||||
return
|
||||
|
||||
idx = selected[0].data(Qt.ItemDataRole.UserRole)
|
||||
try:
|
||||
idx_int = int(idx)
|
||||
except Exception:
|
||||
return
|
||||
if idx_int < 0 or idx_int >= len(self._plot_items):
|
||||
return
|
||||
|
||||
plot = self._plot_items[idx_int]
|
||||
self.selected_plot_title.setText(plot.label)
|
||||
self.selected_plot_path.setText(str(plot.path))
|
||||
|
||||
pix = QPixmap(str(plot.path))
|
||||
if pix.isNull():
|
||||
self.plot_view.clear()
|
||||
return
|
||||
self.plot_view.set_pixmap(pix, fit=True)
|
||||
|
||||
def _infer_run_directories(self, model: Dict[str, Any]) -> List[Path]:
|
||||
dirs: List[Path] = []
|
||||
|
||||
# 1) Infer from model_path: .../<run>/weights/best.pt -> <run>
|
||||
model_path = model.get("model_path")
|
||||
if model_path:
|
||||
try:
|
||||
p = Path(str(model_path)).expanduser()
|
||||
if p.name.lower().endswith(".pt"):
|
||||
# If it lives under weights/, use parent.parent.
|
||||
if p.parent.name == "weights" and p.parent.parent.exists():
|
||||
dirs.append(p.parent.parent)
|
||||
elif p.parent.exists():
|
||||
dirs.append(p.parent)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 2) Look at training_params.stage_results[].results.save_dir
|
||||
training_params = model.get("training_params") or {}
|
||||
stage_results = None
|
||||
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 save_dir:
|
||||
try:
|
||||
save_path = Path(str(save_dir)).expanduser()
|
||||
if save_path.exists():
|
||||
dirs.append(save_path)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
# Deduplicate while preserving order.
|
||||
unique: List[Path] = []
|
||||
seen: set[str] = set()
|
||||
for d in dirs:
|
||||
try:
|
||||
resolved = str(d.resolve())
|
||||
except Exception:
|
||||
resolved = str(d)
|
||||
if resolved not in seen and d.exists() and d.is_dir():
|
||||
seen.add(resolved)
|
||||
unique.append(d)
|
||||
return unique
|
||||
|
||||
def _discover_plot_items(self, directories: Sequence[Path]) -> List[_PlotItem]:
|
||||
# Prefer canonical Ultralytics filenames first, then fall back to any png/jpg.
|
||||
preferred_names = [
|
||||
"results.png",
|
||||
"results.jpg",
|
||||
"confusion_matrix.png",
|
||||
"confusion_matrix_normalized.png",
|
||||
"labels.jpg",
|
||||
"labels.png",
|
||||
"BoxPR_curve.png",
|
||||
"BoxP_curve.png",
|
||||
"BoxR_curve.png",
|
||||
"BoxF1_curve.png",
|
||||
"MaskPR_curve.png",
|
||||
"MaskP_curve.png",
|
||||
"MaskR_curve.png",
|
||||
"MaskF1_curve.png",
|
||||
"val_batch0_pred.jpg",
|
||||
"val_batch0_labels.jpg",
|
||||
]
|
||||
|
||||
found: List[_PlotItem] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
for d in directories:
|
||||
# 1) Preferred
|
||||
for name in preferred_names:
|
||||
p = d / name
|
||||
if p.exists() and p.is_file():
|
||||
key = str(p)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
found.append(_PlotItem(label=f"{name} (from {d.name})", path=p))
|
||||
|
||||
# 2) Curated globs
|
||||
for pattern in ("train_batch*.jpg", "val_batch*.jpg", "*curve*.png"):
|
||||
for p in sorted(d.glob(pattern)):
|
||||
if not p.is_file():
|
||||
continue
|
||||
key = str(p)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
found.append(_PlotItem(label=f"{p.name} (from {d.name})", path=p))
|
||||
|
||||
# 3) Fallback: any top-level png/jpg (excluding weights dir contents)
|
||||
for ext in ("*.png", "*.jpg", "*.jpeg", "*.webp"):
|
||||
for p in sorted(d.glob(ext)):
|
||||
if not p.is_file():
|
||||
continue
|
||||
key = str(p)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
found.append(_PlotItem(label=f"{p.name} (from {d.name})", path=p))
|
||||
|
||||
# Keep list bounded to avoid UI overload for huge runs.
|
||||
return found[:60]
|
||||
|
||||
@@ -18,7 +18,7 @@ from PySide6.QtGui import (
|
||||
QPaintEvent,
|
||||
QPolygonF,
|
||||
)
|
||||
from PySide6.QtCore import Qt, QEvent, Signal, QPoint, QPointF, QRect
|
||||
from PySide6.QtCore import Qt, QEvent, Signal, QPoint, QPointF, QRect, QTimer
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from src.utils.image import Image, ImageLoadError
|
||||
@@ -79,9 +79,7 @@ def rdp(points: List[Tuple[float, float]], epsilon: float) -> List[Tuple[float,
|
||||
return [start, end]
|
||||
|
||||
|
||||
def simplify_polyline(
|
||||
points: List[Tuple[float, float]], epsilon: float
|
||||
) -> List[Tuple[float, float]]:
|
||||
def simplify_polyline(points: List[Tuple[float, float]], epsilon: float) -> List[Tuple[float, float]]:
|
||||
"""
|
||||
Simplify a polyline with RDP while preserving closure semantics.
|
||||
|
||||
@@ -145,6 +143,10 @@ class AnnotationCanvasWidget(QWidget):
|
||||
self.zoom_step = 0.1
|
||||
self.zoom_wheel_step = 0.15
|
||||
|
||||
# Auto-fit behavior (opt-in): when enabled, newly loaded images (and resizes)
|
||||
# will scale to fill the available viewport while preserving aspect ratio.
|
||||
self._auto_fit_to_view: bool = False
|
||||
|
||||
# Drawing / interaction state
|
||||
self.is_drawing = False
|
||||
self.polyline_enabled = False
|
||||
@@ -175,6 +177,35 @@ class AnnotationCanvasWidget(QWidget):
|
||||
|
||||
self._setup_ui()
|
||||
|
||||
def set_auto_fit_to_view(self, enabled: bool):
|
||||
"""Enable/disable automatic zoom-to-fit behavior."""
|
||||
self._auto_fit_to_view = bool(enabled)
|
||||
if self._auto_fit_to_view and self.original_pixmap is not None:
|
||||
QTimer.singleShot(0, self.fit_to_view)
|
||||
|
||||
def fit_to_view(self, padding_px: int = 6):
|
||||
"""Zoom the image so it fits the scroll area's viewport (aspect preserved)."""
|
||||
if self.original_pixmap is None:
|
||||
return
|
||||
|
||||
viewport = self.scroll_area.viewport().size()
|
||||
available_w = max(1, int(viewport.width()) - int(padding_px))
|
||||
available_h = max(1, int(viewport.height()) - int(padding_px))
|
||||
|
||||
img_w = max(1, int(self.original_pixmap.width()))
|
||||
img_h = max(1, int(self.original_pixmap.height()))
|
||||
|
||||
scale_w = available_w / img_w
|
||||
scale_h = available_h / img_h
|
||||
new_scale = min(scale_w, scale_h)
|
||||
new_scale = max(self.zoom_min, min(self.zoom_max, float(new_scale)))
|
||||
|
||||
if abs(new_scale - self.zoom_scale) < 1e-4:
|
||||
return
|
||||
|
||||
self.zoom_scale = new_scale
|
||||
self._apply_zoom()
|
||||
|
||||
def _setup_ui(self):
|
||||
"""Setup user interface."""
|
||||
layout = QVBoxLayout()
|
||||
@@ -187,9 +218,7 @@ class AnnotationCanvasWidget(QWidget):
|
||||
|
||||
self.canvas_label = QLabel("No image loaded")
|
||||
self.canvas_label.setAlignment(Qt.AlignCenter)
|
||||
self.canvas_label.setStyleSheet(
|
||||
"QLabel { background-color: #2b2b2b; color: #888; }"
|
||||
)
|
||||
self.canvas_label.setStyleSheet("QLabel { background-color: #2b2b2b; color: #888; }")
|
||||
self.canvas_label.setScaledContents(False)
|
||||
self.canvas_label.setMouseTracking(True)
|
||||
|
||||
@@ -212,9 +241,18 @@ class AnnotationCanvasWidget(QWidget):
|
||||
self.zoom_scale = 1.0
|
||||
self.clear_annotations()
|
||||
self._display_image()
|
||||
logger.debug(
|
||||
f"Loaded image into annotation canvas: {image.width}x{image.height}"
|
||||
)
|
||||
|
||||
# Defer fit-to-view until the widget has a valid viewport size.
|
||||
if self._auto_fit_to_view:
|
||||
QTimer.singleShot(0, self.fit_to_view)
|
||||
|
||||
logger.debug(f"Loaded image into annotation canvas: {image.width}x{image.height}")
|
||||
|
||||
def resizeEvent(self, event):
|
||||
"""Optionally keep the image fitted when the widget is resized."""
|
||||
super().resizeEvent(event)
|
||||
if self._auto_fit_to_view and self.original_pixmap is not None:
|
||||
QTimer.singleShot(0, self.fit_to_view)
|
||||
|
||||
def clear(self):
|
||||
"""Clear the displayed image and all annotations."""
|
||||
@@ -250,12 +288,10 @@ class AnnotationCanvasWidget(QWidget):
|
||||
# Get image data in a format compatible with Qt
|
||||
if self.current_image.channels in (3, 4):
|
||||
image_data = self.current_image.get_rgb()
|
||||
height, width = image_data.shape[:2]
|
||||
else:
|
||||
image_data = self.current_image.get_grayscale()
|
||||
height, width = image_data.shape
|
||||
image_data = self.current_image.get_qt_rgb()
|
||||
|
||||
image_data = np.ascontiguousarray(image_data)
|
||||
height, width = image_data.shape[:2]
|
||||
bytes_per_line = image_data.strides[0]
|
||||
|
||||
qimage = QImage(
|
||||
@@ -263,7 +299,7 @@ class AnnotationCanvasWidget(QWidget):
|
||||
width,
|
||||
height,
|
||||
bytes_per_line,
|
||||
self.current_image.qtimage_format,
|
||||
QImage.Format_RGBX32FPx4, # self.current_image.qtimage_format,
|
||||
).copy() # Copy so Qt owns the buffer even after numpy array goes out of scope
|
||||
|
||||
self.original_pixmap = QPixmap.fromImage(qimage)
|
||||
@@ -291,22 +327,14 @@ class AnnotationCanvasWidget(QWidget):
|
||||
scaled_width,
|
||||
scaled_height,
|
||||
Qt.KeepAspectRatio,
|
||||
(
|
||||
Qt.SmoothTransformation
|
||||
if self.zoom_scale >= 1.0
|
||||
else Qt.FastTransformation
|
||||
),
|
||||
(Qt.SmoothTransformation if self.zoom_scale >= 1.0 else Qt.FastTransformation),
|
||||
)
|
||||
|
||||
scaled_annotations = self.annotation_pixmap.scaled(
|
||||
scaled_width,
|
||||
scaled_height,
|
||||
Qt.KeepAspectRatio,
|
||||
(
|
||||
Qt.SmoothTransformation
|
||||
if self.zoom_scale >= 1.0
|
||||
else Qt.FastTransformation
|
||||
),
|
||||
(Qt.SmoothTransformation if self.zoom_scale >= 1.0 else Qt.FastTransformation),
|
||||
)
|
||||
|
||||
# Composite image and annotations
|
||||
@@ -392,16 +420,11 @@ class AnnotationCanvasWidget(QWidget):
|
||||
y = (pos.y() - offset_y) / self.zoom_scale
|
||||
|
||||
# Check bounds
|
||||
if (
|
||||
0 <= x < self.original_pixmap.width()
|
||||
and 0 <= y < self.original_pixmap.height()
|
||||
):
|
||||
if 0 <= x < self.original_pixmap.width() and 0 <= y < self.original_pixmap.height():
|
||||
return (int(x), int(y))
|
||||
return None
|
||||
|
||||
def _find_polyline_at(
|
||||
self, img_x: float, img_y: float, threshold_px: float = 5.0
|
||||
) -> Optional[int]:
|
||||
def _find_polyline_at(self, img_x: float, img_y: float, threshold_px: float = 5.0) -> Optional[int]:
|
||||
"""
|
||||
Find index of polyline whose geometry is within threshold_px of (img_x, img_y).
|
||||
Returns the index in self.polylines, or None if none is close enough.
|
||||
@@ -423,9 +446,7 @@ class AnnotationCanvasWidget(QWidget):
|
||||
|
||||
# Precise distance to all segments
|
||||
for (x1, y1), (x2, y2) in zip(polyline[:-1], polyline[1:]):
|
||||
d = perpendicular_distance(
|
||||
(img_x, img_y), (float(x1), float(y1)), (float(x2), float(y2))
|
||||
)
|
||||
d = perpendicular_distance((img_x, img_y), (float(x1), float(y1)), (float(x2), float(y2)))
|
||||
if d < best_dist:
|
||||
best_dist = d
|
||||
best_index = idx
|
||||
@@ -626,11 +647,7 @@ class AnnotationCanvasWidget(QWidget):
|
||||
|
||||
def mouseMoveEvent(self, event: QMouseEvent):
|
||||
"""Handle mouse move events for drawing."""
|
||||
if (
|
||||
not self.is_drawing
|
||||
or not self.polyline_enabled
|
||||
or self.annotation_pixmap is None
|
||||
):
|
||||
if not self.is_drawing or not self.polyline_enabled or self.annotation_pixmap is None:
|
||||
super().mouseMoveEvent(event)
|
||||
return
|
||||
|
||||
@@ -690,15 +707,10 @@ class AnnotationCanvasWidget(QWidget):
|
||||
|
||||
if len(simplified) >= 2:
|
||||
# Store polyline and redraw all annotations
|
||||
self._add_polyline(
|
||||
simplified, self.polyline_pen_color, self.polyline_pen_width
|
||||
)
|
||||
self._add_polyline(simplified, self.polyline_pen_color, self.polyline_pen_width)
|
||||
|
||||
# Convert to normalized coordinates for metadata + signal
|
||||
normalized_stroke = [
|
||||
self._image_to_normalized_coords(int(x), int(y))
|
||||
for (x, y) in simplified
|
||||
]
|
||||
normalized_stroke = [self._image_to_normalized_coords(int(x), int(y)) for (x, y) in simplified]
|
||||
self.all_strokes.append(
|
||||
{
|
||||
"points": normalized_stroke,
|
||||
@@ -711,8 +723,7 @@ class AnnotationCanvasWidget(QWidget):
|
||||
# Emit signal with normalized coordinates
|
||||
self.annotation_drawn.emit(normalized_stroke)
|
||||
logger.debug(
|
||||
f"Completed stroke with {len(simplified)} points "
|
||||
f"(normalized len={len(normalized_stroke)})"
|
||||
f"Completed stroke with {len(simplified)} points " f"(normalized len={len(normalized_stroke)})"
|
||||
)
|
||||
|
||||
self.current_stroke = []
|
||||
@@ -752,9 +763,7 @@ class AnnotationCanvasWidget(QWidget):
|
||||
|
||||
# Store polyline as [y_norm, x_norm] to match DB convention and
|
||||
# the expectations of draw_saved_polyline().
|
||||
normalized_polyline = [
|
||||
[y / img_height, x / img_width] for (x, y) in polyline
|
||||
]
|
||||
normalized_polyline = [[y / img_height, x / img_width] for (x, y) in polyline]
|
||||
|
||||
logger.debug(
|
||||
f"Polyline {idx}: {len(polyline)} points, "
|
||||
@@ -774,7 +783,7 @@ class AnnotationCanvasWidget(QWidget):
|
||||
self,
|
||||
polyline: List[List[float]],
|
||||
color: str,
|
||||
width: int = 3,
|
||||
width: int = 1,
|
||||
annotation_id: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
@@ -812,17 +821,13 @@ class AnnotationCanvasWidget(QWidget):
|
||||
|
||||
# Store and redraw using common pipeline
|
||||
pen_color = QColor(color)
|
||||
pen_color.setAlpha(128) # Add semi-transparency
|
||||
pen_color.setAlpha(255) # Add semi-transparency
|
||||
self._add_polyline(img_coords, pen_color, width, annotation_id=annotation_id)
|
||||
|
||||
# Store in all_strokes for consistency (uses normalized coordinates)
|
||||
self.all_strokes.append(
|
||||
{"points": polyline, "color": color, "alpha": 128, "width": width}
|
||||
)
|
||||
self.all_strokes.append({"points": polyline, "color": color, "alpha": 255, "width": width})
|
||||
|
||||
logger.debug(
|
||||
f"Drew saved polyline with {len(polyline)} points in color {color}"
|
||||
)
|
||||
logger.debug(f"Drew saved polyline with {len(polyline)} points in color {color}")
|
||||
|
||||
def draw_saved_bbox(
|
||||
self,
|
||||
@@ -846,9 +851,7 @@ class AnnotationCanvasWidget(QWidget):
|
||||
return
|
||||
|
||||
if len(bbox) != 4:
|
||||
logger.warning(
|
||||
f"Invalid bounding box format: expected 4 values, got {len(bbox)}"
|
||||
)
|
||||
logger.warning(f"Invalid bounding box format: expected 4 values, got {len(bbox)}")
|
||||
return
|
||||
|
||||
# Convert normalized coordinates to image coordinates (for logging/debug)
|
||||
@@ -869,15 +872,11 @@ class AnnotationCanvasWidget(QWidget):
|
||||
# in _redraw_annotations() together with all polylines.
|
||||
pen_color = QColor(color)
|
||||
pen_color.setAlpha(128) # Add semi-transparency
|
||||
self.bboxes.append(
|
||||
[float(x_min_norm), float(y_min_norm), float(x_max_norm), float(y_max_norm)]
|
||||
)
|
||||
self.bboxes.append([float(x_min_norm), float(y_min_norm), float(x_max_norm), float(y_max_norm)])
|
||||
self.bbox_meta.append({"color": pen_color, "width": int(width), "label": label})
|
||||
|
||||
# Store in all_strokes for consistency
|
||||
self.all_strokes.append(
|
||||
{"bbox": bbox, "color": color, "alpha": 128, "width": width, "label": label}
|
||||
)
|
||||
self.all_strokes.append({"bbox": bbox, "color": color, "alpha": 128, "width": width, "label": label})
|
||||
|
||||
# Redraw overlay (polylines + all bounding boxes)
|
||||
self._redraw_annotations()
|
||||
|
||||
@@ -1,16 +1,21 @@
|
||||
"""
|
||||
YOLO model wrapper for the microscopy object detection application.
|
||||
Provides a clean interface to YOLOv8 for training, validation, and inference.
|
||||
"""YOLO model wrapper for the microscopy object detection application.
|
||||
|
||||
Notes on 16-bit TIFF support:
|
||||
- Ultralytics training defaults assume 8-bit images and normalize by dividing by 255.
|
||||
- This project can patch Ultralytics at runtime to decode TIFFs via `tifffile` and
|
||||
normalize `uint16` correctly.
|
||||
|
||||
See [`apply_ultralytics_16bit_tiff_patches()`](src/utils/ultralytics_16bit_patch.py:1).
|
||||
"""
|
||||
|
||||
from ultralytics import YOLO
|
||||
from pathlib import Path
|
||||
from typing import Optional, List, Dict, Callable, Any
|
||||
import torch
|
||||
import tempfile
|
||||
import os
|
||||
from src.utils.image import Image, convert_grayscale_to_rgb_preserve_range
|
||||
from src.utils.image import Image
|
||||
from src.utils.logger import get_logger
|
||||
from src.utils.ultralytics_16bit_patch import apply_ultralytics_16bit_tiff_patches
|
||||
|
||||
|
||||
logger = get_logger(__name__)
|
||||
@@ -31,6 +36,9 @@ class YOLOWrapper:
|
||||
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
logger.info(f"YOLOWrapper initialized with device: {self.device}")
|
||||
|
||||
# Apply Ultralytics runtime patches early (before first import/instantiation of YOLO datasets/trainers).
|
||||
apply_ultralytics_16bit_tiff_patches()
|
||||
|
||||
def load_model(self) -> bool:
|
||||
"""
|
||||
Load YOLO model from path.
|
||||
@@ -40,6 +48,9 @@ class YOLOWrapper:
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Loading YOLO model from {self.model_path}")
|
||||
# Import YOLO lazily to ensure runtime patches are applied first.
|
||||
from ultralytics import YOLO
|
||||
|
||||
self.model = YOLO(self.model_path)
|
||||
self.model.to(self.device)
|
||||
logger.info("Model loaded successfully")
|
||||
@@ -85,9 +96,17 @@ class YOLOWrapper:
|
||||
|
||||
try:
|
||||
logger.info(f"Starting training: {name}")
|
||||
logger.info(
|
||||
f"Data: {data_yaml}, Epochs: {epochs}, Batch: {batch}, ImgSz: {imgsz}"
|
||||
)
|
||||
logger.info(f"Data: {data_yaml}, Epochs: {epochs}, Batch: {batch}, ImgSz: {imgsz}")
|
||||
|
||||
# Defaults for 16-bit safety: disable augmentations that force uint8 and HSV ops that assume 0..255.
|
||||
# Users can override by passing explicit kwargs.
|
||||
kwargs.setdefault("mosaic", 0.0)
|
||||
kwargs.setdefault("mixup", 0.0)
|
||||
kwargs.setdefault("cutmix", 0.0)
|
||||
kwargs.setdefault("copy_paste", 0.0)
|
||||
kwargs.setdefault("hsv_h", 0.0)
|
||||
kwargs.setdefault("hsv_s", 0.0)
|
||||
kwargs.setdefault("hsv_v", 0.0)
|
||||
|
||||
# Train the model
|
||||
results = self.model.train(
|
||||
@@ -128,9 +147,7 @@ class YOLOWrapper:
|
||||
|
||||
try:
|
||||
logger.info(f"Starting validation on {split} split")
|
||||
results = self.model.val(
|
||||
data=data_yaml, split=split, device=self.device, **kwargs
|
||||
)
|
||||
results = self.model.val(data=data_yaml, split=split, device=self.device, **kwargs)
|
||||
|
||||
logger.info("Validation completed successfully")
|
||||
return self._format_validation_results(results)
|
||||
@@ -169,17 +186,18 @@ class YOLOWrapper:
|
||||
raise RuntimeError(f"Failed to load model from {self.model_path}")
|
||||
|
||||
prepared_source, cleanup_path = self._prepare_source(source)
|
||||
|
||||
imgsz = 1088
|
||||
try:
|
||||
logger.info(f"Running inference on {source}")
|
||||
logger.info(f"Running inference on {source} -> prepared_source {prepared_source}")
|
||||
results = self.model.predict(
|
||||
source=prepared_source,
|
||||
source=source,
|
||||
conf=conf,
|
||||
iou=iou,
|
||||
save=save,
|
||||
save_txt=save_txt,
|
||||
save_conf=save_conf,
|
||||
device=self.device,
|
||||
imgsz=imgsz,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@@ -195,13 +213,9 @@ class YOLOWrapper:
|
||||
try:
|
||||
os.remove(cleanup_path)
|
||||
except OSError as cleanup_error:
|
||||
logger.warning(
|
||||
f"Failed to delete temporary RGB image {cleanup_path}: {cleanup_error}"
|
||||
)
|
||||
logger.warning(f"Failed to delete temporary RGB image {cleanup_path}: {cleanup_error}")
|
||||
|
||||
def export(
|
||||
self, format: str = "onnx", output_path: Optional[str] = None, **kwargs
|
||||
) -> str:
|
||||
def export(self, format: str = "onnx", output_path: Optional[str] = None, **kwargs) -> str:
|
||||
"""
|
||||
Export model to different format.
|
||||
|
||||
@@ -236,21 +250,13 @@ class YOLOWrapper:
|
||||
if source_path.is_file():
|
||||
try:
|
||||
img_obj = Image(source_path)
|
||||
pil_img = img_obj.pil_image
|
||||
if len(pil_img.getbands()) == 1:
|
||||
rgb_img = convert_grayscale_to_rgb_preserve_range(pil_img)
|
||||
else:
|
||||
rgb_img = pil_img.convert("RGB")
|
||||
|
||||
suffix = source_path.suffix or ".png"
|
||||
tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
|
||||
tmp_path = tmp.name
|
||||
tmp.close()
|
||||
rgb_img.save(tmp_path)
|
||||
img_obj.save(tmp_path)
|
||||
cleanup_path = tmp_path
|
||||
logger.info(
|
||||
f"Converted image {source_path} to RGB for inference at {tmp_path}"
|
||||
)
|
||||
logger.info(f"Converted image {source_path} to RGB for inference at {tmp_path}")
|
||||
return tmp_path, cleanup_path
|
||||
except Exception as convert_error:
|
||||
logger.warning(
|
||||
@@ -263,9 +269,7 @@ class YOLOWrapper:
|
||||
"""Format training results into dictionary."""
|
||||
try:
|
||||
# Get the results dict
|
||||
results_dict = (
|
||||
results.results_dict if hasattr(results, "results_dict") else {}
|
||||
)
|
||||
results_dict = results.results_dict if hasattr(results, "results_dict") else {}
|
||||
|
||||
formatted = {
|
||||
"success": True,
|
||||
@@ -298,9 +302,7 @@ class YOLOWrapper:
|
||||
"mAP50-95": float(box_metrics.map),
|
||||
"precision": float(box_metrics.mp),
|
||||
"recall": float(box_metrics.mr),
|
||||
"fitness": (
|
||||
float(results.fitness) if hasattr(results, "fitness") else 0.0
|
||||
),
|
||||
"fitness": (float(results.fitness) if hasattr(results, "fitness") else 0.0),
|
||||
}
|
||||
|
||||
# Add per-class metrics if available
|
||||
@@ -310,11 +312,7 @@ class YOLOWrapper:
|
||||
if idx < len(box_metrics.ap):
|
||||
class_metrics[name] = {
|
||||
"ap": float(box_metrics.ap[idx]),
|
||||
"ap50": (
|
||||
float(box_metrics.ap50[idx])
|
||||
if hasattr(box_metrics, "ap50")
|
||||
else 0.0
|
||||
),
|
||||
"ap50": (float(box_metrics.ap50[idx]) if hasattr(box_metrics, "ap50") else 0.0),
|
||||
}
|
||||
formatted["class_metrics"] = class_metrics
|
||||
|
||||
@@ -347,21 +345,15 @@ class YOLOWrapper:
|
||||
"class_id": int(boxes.cls[i]),
|
||||
"class_name": result.names[int(boxes.cls[i])],
|
||||
"confidence": float(boxes.conf[i]),
|
||||
"bbox_normalized": [
|
||||
float(v) for v in xyxyn
|
||||
], # [x_min, y_min, x_max, y_max]
|
||||
"bbox_absolute": [
|
||||
float(v) for v in boxes.xyxy[i].cpu().numpy()
|
||||
], # Absolute pixels
|
||||
"bbox_normalized": [float(v) for v in xyxyn], # [x_min, y_min, x_max, y_max]
|
||||
"bbox_absolute": [float(v) for v in boxes.xyxy[i].cpu().numpy()], # Absolute pixels
|
||||
}
|
||||
|
||||
# Extract segmentation mask if available
|
||||
if has_masks:
|
||||
try:
|
||||
# Get the mask for this detection
|
||||
mask_data = result.masks.xy[
|
||||
i
|
||||
] # Polygon coordinates in absolute pixels
|
||||
mask_data = result.masks.xy[i] # Polygon coordinates in absolute pixels
|
||||
|
||||
# Convert to normalized coordinates
|
||||
if len(mask_data) > 0:
|
||||
@@ -374,9 +366,7 @@ class YOLOWrapper:
|
||||
else:
|
||||
detection["segmentation_mask"] = None
|
||||
except Exception as mask_error:
|
||||
logger.warning(
|
||||
f"Error extracting mask for detection {i}: {mask_error}"
|
||||
)
|
||||
logger.warning(f"Error extracting mask for detection {i}: {mask_error}")
|
||||
detection["segmentation_mask"] = None
|
||||
else:
|
||||
detection["segmentation_mask"] = None
|
||||
@@ -390,9 +380,7 @@ class YOLOWrapper:
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def convert_bbox_format(
|
||||
bbox: List[float], format_from: str = "xywh", format_to: str = "xyxy"
|
||||
) -> List[float]:
|
||||
def convert_bbox_format(bbox: List[float], format_from: str = "xywh", format_to: str = "xyxy") -> List[float]:
|
||||
"""
|
||||
Convert bounding box between formats.
|
||||
|
||||
|
||||
@@ -54,7 +54,7 @@ class ConfigManager:
|
||||
"models_directory": "data/models",
|
||||
"base_model_choices": [
|
||||
"yolov8s-seg.pt",
|
||||
"yolov11s-seg.pt",
|
||||
"yolo11s-seg.pt",
|
||||
],
|
||||
},
|
||||
"training": {
|
||||
@@ -225,6 +225,4 @@ class ConfigManager:
|
||||
|
||||
def get_allowed_extensions(self) -> list:
|
||||
"""Get list of allowed image file extensions."""
|
||||
return self.get(
|
||||
"image_repository.allowed_extensions", Image.SUPPORTED_EXTENSIONS
|
||||
)
|
||||
return self.get("image_repository.allowed_extensions", Image.SUPPORTED_EXTENSIONS)
|
||||
|
||||
103
src/utils/create_mask_from_detection.py
Normal file
103
src/utils/create_mask_from_detection.py
Normal file
@@ -0,0 +1,103 @@
|
||||
import numpy as np
|
||||
|
||||
from pathlib import Path
|
||||
from skimage.draw import polygon
|
||||
from tifffile import TiffFile
|
||||
|
||||
from src.database.db_manager import DatabaseManager
|
||||
|
||||
|
||||
def read_image(image_path: Path) -> np.ndarray:
|
||||
metadata = {}
|
||||
with TiffFile(image_path) as tif:
|
||||
image = tif.asarray()
|
||||
metadata = tif.imagej_metadata
|
||||
return image, metadata
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
polygon_vertices = np.array([[10, 10], [50, 10], [50, 50], [10, 50]])
|
||||
image = np.zeros((100, 100), dtype=np.uint8)
|
||||
rr, cc = polygon(polygon_vertices[:, 0], polygon_vertices[:, 1])
|
||||
image[rr, cc] = 255
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
db = DatabaseManager()
|
||||
model_name = "c17"
|
||||
model_id = db.get_models(filters={"model_name": model_name})[0]["id"]
|
||||
print(f"Model name {model_name}, id {model_id}")
|
||||
detections = db.get_detections(filters={"model_id": model_id})
|
||||
|
||||
file_stems = set()
|
||||
|
||||
for detection in detections:
|
||||
file_stems.add(detection["image_filename"].split("_")[0])
|
||||
|
||||
print("Files:", file_stems)
|
||||
|
||||
for stem in file_stems:
|
||||
print(stem)
|
||||
detections = db.get_detections(filters={"model_id": model_id, "i.filename": f"LIKE %{stem}%"})
|
||||
annotations = []
|
||||
for detection in detections:
|
||||
source_path = Path(detection["metadata"]["source_path"])
|
||||
image, metadata = read_image(source_path)
|
||||
|
||||
offset = np.array(list(map(int, metadata["tile_section"].split(","))))[::-1]
|
||||
scale = np.array(list(map(int, metadata["patch_size"].split(","))))[::-1]
|
||||
# tile_size = np.array(list(map(int, metadata["tile_size"].split(","))))
|
||||
segmentation = np.array(detection["segmentation_mask"]) # * tile_size
|
||||
|
||||
# print(source_path, image, metadata, segmentation.shape)
|
||||
# print(offset)
|
||||
# print(scale)
|
||||
# print(segmentation)
|
||||
|
||||
# segmentation = (segmentation + offset * tile_size) / (tile_size * scale)
|
||||
segmentation = (segmentation + offset) / scale
|
||||
|
||||
yolo_annotation = f"{detection['metadata']['class_id']} " + " ".join(
|
||||
[f"{x:.6f} {y:.6f}" for x, y in segmentation]
|
||||
)
|
||||
annotations.append(yolo_annotation)
|
||||
# print(segmentation)
|
||||
# print(yolo_annotation)
|
||||
|
||||
# aa
|
||||
print(
|
||||
" ",
|
||||
detection["model_name"],
|
||||
detection["image_id"],
|
||||
detection["image_filename"],
|
||||
source_path,
|
||||
metadata["label_path"],
|
||||
)
|
||||
# section_i_section_j = detection["image_filename"].split("_")[1].split(".")[0]
|
||||
# print(" ", section_i_section_j)
|
||||
|
||||
label_path = metadata["label_path"]
|
||||
print(" ", label_path)
|
||||
with open(label_path, "w") as f:
|
||||
f.write("\n".join(annotations))
|
||||
|
||||
exit()
|
||||
|
||||
for detection in detections:
|
||||
print(detection["model_name"], detection["image_id"], detection["image_filename"])
|
||||
|
||||
print(detections[0])
|
||||
# polygon_vertices = np.array([[10, 10], [50, 10], [50, 50], [10, 50]])
|
||||
|
||||
# image = np.zeros((100, 100), dtype=np.uint8)
|
||||
|
||||
# rr, cc = polygon(polygon_vertices[:, 0], polygon_vertices[:, 1])
|
||||
|
||||
# image[rr, cc] = 255
|
||||
|
||||
# import matplotlib.pyplot as plt
|
||||
|
||||
# plt.imshow(image, cmap='gray')
|
||||
# plt.show()
|
||||
@@ -6,16 +6,55 @@ import cv2
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple, Union
|
||||
from PIL import Image as PILImage
|
||||
|
||||
from src.utils.logger import get_logger
|
||||
from src.utils.file_utils import validate_file_path, is_image_file
|
||||
|
||||
from PySide6.QtGui import QImage
|
||||
|
||||
from tifffile import imread, imwrite
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def get_pseudo_rgb(arr: np.ndarray, gamma: float = 0.5) -> np.ndarray:
|
||||
"""
|
||||
Convert a grayscale image to a pseudo-RGB image using a gamma correction.
|
||||
|
||||
Args:
|
||||
arr: Input grayscale image as numpy array
|
||||
|
||||
Returns:
|
||||
Pseudo-RGB image as numpy array
|
||||
"""
|
||||
if arr.ndim != 2:
|
||||
raise ValueError("Input array must be a grayscale image with shape (H, W)")
|
||||
|
||||
a1 = arr.copy().astype(np.float32)
|
||||
a1 -= np.percentile(a1, 2)
|
||||
a1[a1 < 0] = 0
|
||||
p999 = np.percentile(a1, 99.9)
|
||||
a1[a1 > p999] = p999
|
||||
a1 /= a1.max()
|
||||
|
||||
if 1:
|
||||
a2 = a1.copy()
|
||||
a2 = a2**gamma
|
||||
a2 /= a2.max()
|
||||
|
||||
a3 = a1.copy()
|
||||
p9999 = np.percentile(a3, 99.99)
|
||||
a3[a3 > p9999] = p9999
|
||||
a3 /= a3.max()
|
||||
|
||||
# return np.stack([a1, np.zeros(a1.shape), np.zeros(a1.shape)], axis=0)
|
||||
# return np.stack([a2, np.zeros(a1.shape), np.zeros(a1.shape)], axis=0)
|
||||
out = np.stack([a1, a2, a3], axis=0)
|
||||
# print(any(np.isnan(out).flatten()))
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ImageLoadError(Exception):
|
||||
"""Exception raised when an image cannot be loaded."""
|
||||
|
||||
@@ -54,7 +93,6 @@ class Image:
|
||||
"""
|
||||
self.path = Path(image_path)
|
||||
self._data: Optional[np.ndarray] = None
|
||||
self._pil_image: Optional[PILImage.Image] = None
|
||||
self._width: int = 0
|
||||
self._height: int = 0
|
||||
self._channels: int = 0
|
||||
@@ -80,11 +118,14 @@ class Image:
|
||||
if not is_image_file(str(self.path), self.SUPPORTED_EXTENSIONS):
|
||||
ext = self.path.suffix.lower()
|
||||
raise ImageLoadError(
|
||||
f"Unsupported image format: {ext}. "
|
||||
f"Supported formats: {', '.join(self.SUPPORTED_EXTENSIONS)}"
|
||||
f"Unsupported image format: {ext}. " f"Supported formats: {', '.join(self.SUPPORTED_EXTENSIONS)}"
|
||||
)
|
||||
|
||||
try:
|
||||
if self.path.suffix.lower() in [".tif", ".tiff"]:
|
||||
self._data = imread(str(self.path))
|
||||
else:
|
||||
# raise NotImplementedError("RGB is not implemented")
|
||||
# Load with OpenCV (returns BGR format)
|
||||
self._data = cv2.imread(str(self.path), cv2.IMREAD_UNCHANGED)
|
||||
|
||||
@@ -92,23 +133,19 @@ class Image:
|
||||
raise ImageLoadError(f"Failed to load image with OpenCV: {self.path}")
|
||||
|
||||
# Extract metadata
|
||||
# print(self._data.shape)
|
||||
if len(self._data.shape) == 2:
|
||||
self._height, self._width = self._data.shape[:2]
|
||||
self._channels = self._data.shape[2] if len(self._data.shape) == 3 else 1
|
||||
self._channels = 1
|
||||
else:
|
||||
self._height, self._width = self._data.shape[1:]
|
||||
self._channels = self._data.shape[0]
|
||||
# self._channels = self._data.shape[2] if len(self._data.shape) == 3 else 1
|
||||
self._format = self.path.suffix.lower().lstrip(".")
|
||||
self._size_bytes = self.path.stat().st_size
|
||||
self._dtype = self._data.dtype
|
||||
|
||||
# Load PIL version for compatibility (convert BGR to RGB)
|
||||
if self._channels == 3:
|
||||
rgb_data = cv2.cvtColor(self._data, cv2.COLOR_BGR2RGB)
|
||||
self._pil_image = PILImage.fromarray(rgb_data)
|
||||
elif self._channels == 4:
|
||||
rgba_data = cv2.cvtColor(self._data, cv2.COLOR_BGRA2RGBA)
|
||||
self._pil_image = PILImage.fromarray(rgba_data)
|
||||
else:
|
||||
# Grayscale
|
||||
self._pil_image = PILImage.fromarray(self._data)
|
||||
|
||||
if 0:
|
||||
logger.info(
|
||||
f"Successfully loaded image: {self.path.name} "
|
||||
f"({self._width}x{self._height}, {self._channels} channels, "
|
||||
@@ -131,18 +168,6 @@ class Image:
|
||||
raise ImageLoadError("Image data not available")
|
||||
return self._data
|
||||
|
||||
@property
|
||||
def pil_image(self) -> PILImage.Image:
|
||||
"""
|
||||
Get image data as PIL Image (RGB or grayscale).
|
||||
|
||||
Returns:
|
||||
PIL Image object
|
||||
"""
|
||||
if self._pil_image is None:
|
||||
raise ImageLoadError("PIL image not available")
|
||||
return self._pil_image
|
||||
|
||||
@property
|
||||
def width(self) -> int:
|
||||
"""Get image width in pixels."""
|
||||
@@ -187,6 +212,7 @@ class Image:
|
||||
@property
|
||||
def dtype(self) -> np.dtype:
|
||||
"""Get the data type of the image array."""
|
||||
|
||||
if self._dtype is None:
|
||||
raise ImageLoadError("Image dtype not available")
|
||||
return self._dtype
|
||||
@@ -206,8 +232,10 @@ class Image:
|
||||
elif self._channels == 1:
|
||||
if self._dtype == np.uint16:
|
||||
return QImage.Format_Grayscale16
|
||||
else:
|
||||
elif self._dtype == np.uint8:
|
||||
return QImage.Format_Grayscale8
|
||||
elif self._dtype == np.float32:
|
||||
return QImage.Format_BGR30
|
||||
else:
|
||||
raise ImageLoadError(f"Unsupported number of channels: {self._channels}")
|
||||
|
||||
@@ -218,12 +246,36 @@ class Image:
|
||||
Returns:
|
||||
Image data in RGB format as numpy array
|
||||
"""
|
||||
if self._channels == 3:
|
||||
return cv2.cvtColor(self._data, cv2.COLOR_BGR2RGB)
|
||||
if self.channels == 1:
|
||||
img = get_pseudo_rgb(self.data)
|
||||
self._dtype = img.dtype
|
||||
return img, True
|
||||
|
||||
elif self._channels == 3:
|
||||
return cv2.cvtColor(self._data, cv2.COLOR_BGR2RGB), False
|
||||
elif self._channels == 4:
|
||||
return cv2.cvtColor(self._data, cv2.COLOR_BGRA2RGBA)
|
||||
return cv2.cvtColor(self._data, cv2.COLOR_BGRA2RGBA), False
|
||||
|
||||
else:
|
||||
return self._data
|
||||
raise NotImplementedError
|
||||
|
||||
# else:
|
||||
# return self._data
|
||||
|
||||
def get_qt_rgb(self) -> np.ascontiguousarray:
|
||||
# we keep data as (C, H, W)
|
||||
_img, pseudo = self.get_rgb()
|
||||
|
||||
if pseudo:
|
||||
img = np.zeros((self.height, self.width, 4), dtype=np.float32)
|
||||
img[..., 0] = _img[0] # R gradient
|
||||
img[..., 1] = _img[1] # G gradient
|
||||
img[..., 2] = _img[2] # B constant
|
||||
img[..., 3] = 1.0 # A = 1.0 (opaque)
|
||||
|
||||
return np.ascontiguousarray(img)
|
||||
else:
|
||||
return np.ascontiguousarray(_img)
|
||||
|
||||
def get_grayscale(self) -> np.ndarray:
|
||||
"""
|
||||
@@ -277,11 +329,26 @@ class Image:
|
||||
"""
|
||||
return self._channels >= 3
|
||||
|
||||
def save(self, path: Union[str, Path], pseudo_rgb: bool = True) -> None:
|
||||
|
||||
if self.channels == 1:
|
||||
if pseudo_rgb:
|
||||
img = get_pseudo_rgb(self.data)
|
||||
print("Image.save", img.shape)
|
||||
else:
|
||||
img = np.repeat(self.data, 3, axis=2)
|
||||
|
||||
else:
|
||||
raise NotImplementedError("Only grayscale images are supported for now.")
|
||||
|
||||
imwrite(path, data=img)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""String representation of the Image object."""
|
||||
return (
|
||||
f"Image(path='{self.path.name}', "
|
||||
f"shape=({self._width}x{self._height}x{self._channels}), "
|
||||
# Display as HxWxC to match the conventional NumPy shape semantics.
|
||||
f"shape=({self._height}x{self._width}x{self._channels}), "
|
||||
f"format={self._format}, "
|
||||
f"size={self.size_mb:.2f}MB)"
|
||||
)
|
||||
@@ -291,38 +358,13 @@ class Image:
|
||||
return self.__repr__()
|
||||
|
||||
|
||||
def convert_grayscale_to_rgb_preserve_range(
|
||||
pil_image: PILImage.Image,
|
||||
) -> PILImage.Image:
|
||||
"""Convert a single-channel PIL image to RGB while preserving dynamic range.
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
Args:
|
||||
pil_image: Single-channel PIL image (e.g., 16-bit grayscale).
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--path", type=str, required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
Returns:
|
||||
PIL Image in RGB mode with intensities normalized to 0-255.
|
||||
"""
|
||||
|
||||
if pil_image.mode == "RGB":
|
||||
return pil_image
|
||||
|
||||
grayscale = np.array(pil_image)
|
||||
if grayscale.ndim == 3:
|
||||
grayscale = grayscale[:, :, 0]
|
||||
|
||||
original_dtype = grayscale.dtype
|
||||
grayscale = grayscale.astype(np.float32)
|
||||
|
||||
if grayscale.size == 0:
|
||||
return PILImage.new("RGB", pil_image.size, color=(0, 0, 0))
|
||||
|
||||
if np.issubdtype(original_dtype, np.integer):
|
||||
denom = float(max(np.iinfo(original_dtype).max, 1))
|
||||
else:
|
||||
max_val = float(grayscale.max())
|
||||
denom = max(max_val, 1.0)
|
||||
|
||||
grayscale = np.clip(grayscale / denom, 0.0, 1.0)
|
||||
grayscale_u8 = (grayscale * 255.0).round().astype(np.uint8)
|
||||
rgb_arr = np.repeat(grayscale_u8[:, :, None], 3, axis=2)
|
||||
return PILImage.fromarray(rgb_arr, mode="RGB")
|
||||
img = Image(args.path)
|
||||
img.save(args.path + "test.tif")
|
||||
print(img)
|
||||
|
||||
@@ -12,23 +12,38 @@ class UT:
|
||||
Operetta files along with rois drawn in ImageJ
|
||||
"""
|
||||
|
||||
def __init__(self, roifile_fn: Path):
|
||||
def __init__(self, roifile_fn: Path, no_labels: bool):
|
||||
self.roifile_fn = roifile_fn
|
||||
print("is file", self.roifile_fn.is_file())
|
||||
self.rois = None
|
||||
if no_labels:
|
||||
self.rois = ImagejRoi.fromfile(self.roifile_fn)
|
||||
self.stem = self.roifile_fn.stem.strip("-RoiSet")
|
||||
print(self.roifile_fn.stem)
|
||||
print(self.roifile_fn.parent.parts[-1])
|
||||
if "Roi-" in self.roifile_fn.stem:
|
||||
self.stem = self.roifile_fn.stem.split("Roi-")[1]
|
||||
else:
|
||||
self.stem = self.roifile_fn.parent.parts[-1]
|
||||
|
||||
else:
|
||||
self.roifile_fn = roifile_fn / roifile_fn.parts[-1]
|
||||
self.stem = self.roifile_fn.stem
|
||||
|
||||
print(self.roifile_fn)
|
||||
|
||||
print(self.stem)
|
||||
self.image, self.image_props = self._load_images()
|
||||
|
||||
def _load_images(self):
|
||||
"""Loading sequence of tif files
|
||||
array sequence is CZYX
|
||||
"""
|
||||
print(self.roifile_fn.parent, self.stem)
|
||||
fns = list(self.roifile_fn.parent.glob(f"{self.stem}*.tif*"))
|
||||
print("Loading images:", self.roifile_fn.parent, self.stem)
|
||||
fns = list(self.roifile_fn.parent.glob(f"{self.stem.lower()}*.tif*"))
|
||||
stems = [fn.stem.split(self.stem)[-1] for fn in fns]
|
||||
n_ch = len(set([stem.split("-ch")[-1].split("t")[0] for stem in stems]))
|
||||
n_p = len(set([stem.split("-")[0] for stem in stems]))
|
||||
n_t = len(set([stem.split("t")[1] for stem in stems]))
|
||||
print(n_ch, n_p, n_t)
|
||||
|
||||
with TiffFile(fns[0]) as tif:
|
||||
img = tif.asarray()
|
||||
@@ -42,6 +57,7 @@ class UT:
|
||||
"height": h,
|
||||
"dtype": dtype,
|
||||
}
|
||||
print("Image props", self.image_props)
|
||||
|
||||
image_stack = np.zeros((n_ch, n_p, w, h), dtype=dtype)
|
||||
for fn in fns:
|
||||
@@ -49,7 +65,7 @@ class UT:
|
||||
img = tif.asarray()
|
||||
stem = fn.stem.split(self.stem)[-1]
|
||||
ch = int(stem.split("-ch")[-1].split("t")[0])
|
||||
p = int(stem.split("-")[0].lstrip("p"))
|
||||
p = int(stem.split("-")[0].split("p")[1])
|
||||
t = int(stem.split("t")[1])
|
||||
print(fn.stem, "ch", ch, "p", p, "t", t)
|
||||
image_stack[ch - 1, p - 1] = img
|
||||
@@ -82,10 +98,19 @@ class UT:
|
||||
):
|
||||
"""Export rois to a file"""
|
||||
with open(path / subfolder / f"{self.stem}.txt", "w") as f:
|
||||
for roi in self.rois:
|
||||
# TODO add image coordinates normalization
|
||||
coords = ""
|
||||
for x, y in roi.subpixel_coordinates:
|
||||
for i, roi in enumerate(self.rois):
|
||||
rc = roi.subpixel_coordinates
|
||||
if rc is None:
|
||||
print(f"No coordinates: {self.roifile_fn}, element {i}, out of {len(self.rois)}")
|
||||
continue
|
||||
xmn, ymn = rc.min(axis=0)
|
||||
xmx, ymx = rc.max(axis=0)
|
||||
xc = (xmn + xmx) / 2
|
||||
yc = (ymn + ymx) / 2
|
||||
bw = xmx - xmn
|
||||
bh = ymx - ymn
|
||||
coords = f"{xc/self.width} {yc/self.height} {bw/self.width} {bh/self.height} "
|
||||
for x, y in rc:
|
||||
coords += f"{x/self.width} {y/self.height} "
|
||||
f.write(f"{class_index} {coords}\n")
|
||||
|
||||
@@ -104,6 +129,7 @@ class UT:
|
||||
self.image = np.max(self.image[channel], axis=0)
|
||||
print(self.image.shape)
|
||||
|
||||
print(path / subfolder / f"{self.stem}.tif")
|
||||
with TiffWriter(path / subfolder / f"{self.stem}.tif") as tif:
|
||||
tif.write(self.image)
|
||||
|
||||
@@ -112,11 +138,31 @@ if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("input", type=Path)
|
||||
parser.add_argument("output", type=Path)
|
||||
parser.add_argument("-i", "--input", nargs="*", type=Path)
|
||||
parser.add_argument("-o", "--output", type=Path)
|
||||
parser.add_argument(
|
||||
"--no-labels",
|
||||
action="store_false",
|
||||
help="Source does not have labels, export only images",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
for rfn in args.input.glob("*.zip"):
|
||||
ut = UT(rfn)
|
||||
# print(args)
|
||||
# aa
|
||||
|
||||
for path in args.input:
|
||||
print("Path:", path)
|
||||
if not args.no_labels:
|
||||
print("No labels")
|
||||
ut = UT(path, args.no_labels)
|
||||
ut.export_image(args.output, plane_mode="max projection", channel=0)
|
||||
|
||||
else:
|
||||
for rfn in Path(path).glob("*.zip"):
|
||||
# if Path(path).suffix == ".zip":
|
||||
print("Roi FN:", rfn)
|
||||
ut = UT(rfn, args.no_labels)
|
||||
ut.export_rois(args.output, class_index=0)
|
||||
ut.export_image(args.output, plane_mode="max projection", channel=0)
|
||||
|
||||
print()
|
||||
|
||||
368
src/utils/image_splitter.py
Normal file
368
src/utils/image_splitter.py
Normal file
@@ -0,0 +1,368 @@
|
||||
import numpy as np
|
||||
|
||||
from pathlib import Path
|
||||
from tifffile import imread, imwrite
|
||||
from shapely.geometry import LineString
|
||||
from copy import deepcopy
|
||||
from scipy.ndimage import zoom
|
||||
|
||||
|
||||
# debug
|
||||
from src.utils.image import Image
|
||||
from show_yolo_seg import draw_annotations
|
||||
|
||||
import pylab as plt
|
||||
import cv2
|
||||
|
||||
|
||||
class Label:
|
||||
def __init__(self, yolo_annotation: str):
|
||||
class_id, bbox, polygon = self.parse_yolo_annotation(yolo_annotation)
|
||||
self.class_id = class_id
|
||||
self.bbox = bbox
|
||||
self.polygon = polygon
|
||||
|
||||
def parse_yolo_annotation(self, yolo_annotation: str):
|
||||
class_id, *coords = yolo_annotation.split()
|
||||
class_id = int(class_id)
|
||||
bbox = np.array(coords[:4], dtype=np.float32)
|
||||
polygon = np.array(coords[4:], dtype=np.float32).reshape(-1, 2) if len(coords) > 4 else None
|
||||
if not any(np.isclose(polygon[0], polygon[-1])):
|
||||
polygon = np.vstack([polygon, polygon[0]])
|
||||
return class_id, bbox, polygon
|
||||
|
||||
def offset_label(
|
||||
self,
|
||||
img_w,
|
||||
img_h,
|
||||
distance: float = 1.0,
|
||||
cap_style: int = 2,
|
||||
join_style: int = 2,
|
||||
):
|
||||
if self.polygon is None:
|
||||
self.bbox = np.array(
|
||||
[
|
||||
self.bbox[0] - distance if self.bbox[0] - distance > 0 else 0,
|
||||
self.bbox[1] - distance if self.bbox[1] - distance > 0 else 0,
|
||||
self.bbox[2] + distance if self.bbox[2] + distance < 1 else 1,
|
||||
self.bbox[3] + distance if self.bbox[3] + distance < 1 else 1,
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
return self.bbox
|
||||
|
||||
def coords_are_normalized(coords):
|
||||
# If every coordinate is between 0 and 1 (inclusive-ish), assume normalized
|
||||
print(coords)
|
||||
# if not coords:
|
||||
# return False
|
||||
return all(max(coords.flatten)) <= 1.001
|
||||
|
||||
def poly_to_pts(coords, img_w, img_h):
|
||||
# coords: [x1 y1 x2 y2 ...] either normalized or absolute
|
||||
# if coords_are_normalized(coords):
|
||||
coords = [coords[i] * (img_w if i % 2 == 0 else img_h) for i in range(len(coords))]
|
||||
pts = np.array(coords, dtype=np.int32).reshape(-1, 2)
|
||||
return pts
|
||||
|
||||
pts = poly_to_pts(self.polygon, img_w, img_h)
|
||||
line = LineString(pts)
|
||||
# Buffer distance in pixels
|
||||
buffered = line.buffer(distance=distance, cap_style=cap_style, join_style=join_style)
|
||||
self.polygon = np.array(buffered.exterior.coords, dtype=np.float32) / (img_w, img_h)
|
||||
xmn, ymn = self.polygon.min(axis=0)
|
||||
xmx, ymx = self.polygon.max(axis=0)
|
||||
xc = (xmn + xmx) / 2
|
||||
yc = (ymn + ymx) / 2
|
||||
bw = xmx - xmn
|
||||
bh = ymx - ymn
|
||||
self.bbox = np.array([xc, yc, bw, bh], dtype=np.float32)
|
||||
|
||||
return self.bbox, self.polygon
|
||||
|
||||
def translate(self, x, y, scale_x, scale_y):
|
||||
self.bbox[0] -= x
|
||||
self.bbox[0] *= scale_x
|
||||
self.bbox[1] -= y
|
||||
self.bbox[1] *= scale_y
|
||||
self.bbox[2] *= scale_x
|
||||
self.bbox[3] *= scale_y
|
||||
if self.polygon is not None:
|
||||
self.polygon[:, 0] -= x
|
||||
self.polygon[:, 0] *= scale_x
|
||||
self.polygon[:, 1] -= y
|
||||
self.polygon[:, 1] *= scale_y
|
||||
|
||||
def in_range(self, hrange, wrange):
|
||||
xc, yc, h, w = self.bbox
|
||||
x1 = xc - w / 2
|
||||
y1 = yc - h / 2
|
||||
x2 = xc + w / 2
|
||||
y2 = yc + h / 2
|
||||
truth_val = (
|
||||
xc >= wrange[0]
|
||||
and x1 <= wrange[1]
|
||||
and x2 >= wrange[0]
|
||||
and x2 <= wrange[1]
|
||||
and y1 >= hrange[0]
|
||||
and y1 <= hrange[1]
|
||||
and y2 >= hrange[0]
|
||||
and y2 <= hrange[1]
|
||||
)
|
||||
|
||||
print(x1, x2, wrange, y1, y2, hrange, truth_val)
|
||||
return truth_val
|
||||
|
||||
def to_string(self, bbox: list = None, polygon: list = None):
|
||||
coords = ""
|
||||
if bbox is None:
|
||||
bbox = self.bbox
|
||||
# coords += " ".join([f"{x:.6f}" for x in self.bbox])
|
||||
if polygon is None:
|
||||
polygon = self.polygon
|
||||
if self.polygon is not None:
|
||||
coords += " " + " ".join([f"{x:.6f} {y:.6f}" for x, y in self.polygon])
|
||||
return f"{self.class_id} {coords}"
|
||||
|
||||
def __str__(self):
|
||||
return f"Class: {self.class_id}, BBox: {self.bbox}, Polygon: {self.polygon}"
|
||||
|
||||
|
||||
class YoloLabelReader:
|
||||
def __init__(self, label_path: Path):
|
||||
self.label_path = label_path
|
||||
self.labels = self._read_labels()
|
||||
|
||||
def _read_labels(self):
|
||||
with open(self.label_path, "r") as f:
|
||||
labels = [Label(line) for line in f.readlines()]
|
||||
|
||||
return labels
|
||||
|
||||
def get_labels(self, hrange, wrange):
|
||||
"""hrange and wrange are tuples of (start, end) normalized to [0, 1]"""
|
||||
labels = []
|
||||
# print(hrange, wrange)
|
||||
for lbl in self.labels:
|
||||
# print(lbl)
|
||||
if lbl.in_range(hrange, wrange):
|
||||
labels.append(lbl)
|
||||
return labels if len(labels) > 0 else None
|
||||
|
||||
def __get_item__(self, index):
|
||||
return self.labels[index]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.labels)
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self.labels)
|
||||
|
||||
|
||||
class ImageSplitter:
|
||||
def __init__(self, image_path: Path, label_path: Path):
|
||||
self.image = imread(image_path)
|
||||
self.image_path = image_path
|
||||
self.label_path = label_path
|
||||
if not label_path.exists():
|
||||
print(f"Label file {label_path} not found")
|
||||
self.labels = None
|
||||
else:
|
||||
self.labels = YoloLabelReader(label_path)
|
||||
|
||||
def split_into_tiles(self, patch_size: tuple = (2, 2)):
|
||||
"""Split image into patches of size patch_size"""
|
||||
hstep, wstep = (
|
||||
self.image.shape[0] // patch_size[0],
|
||||
self.image.shape[1] // patch_size[1],
|
||||
)
|
||||
h, w = self.image.shape[:2]
|
||||
|
||||
for i in range(patch_size[0]):
|
||||
for j in range(patch_size[1]):
|
||||
metadata = {
|
||||
"image_path": str(self.image_path),
|
||||
"label_path": str(self.label_path),
|
||||
"tile_section": f"{i}, {j}",
|
||||
"tile_size": f"{hstep}, {wstep}",
|
||||
"patch_size": f"{patch_size[0]}, {patch_size[1]}",
|
||||
}
|
||||
tile_reference = f"i{i}j{j}"
|
||||
hrange = (i * hstep / h, (i + 1) * hstep / h)
|
||||
wrange = (j * wstep / w, (j + 1) * wstep / w)
|
||||
tile = self.image[i * hstep : (i + 1) * hstep, j * wstep : (j + 1) * wstep]
|
||||
|
||||
labels = None
|
||||
if self.labels is not None:
|
||||
labels = deepcopy(self.labels.get_labels(hrange, wrange))
|
||||
print(id(labels))
|
||||
|
||||
if labels is not None:
|
||||
print(hrange[0], wrange[0])
|
||||
for l in labels:
|
||||
print(l.bbox)
|
||||
[l.translate(wrange[0], hrange[0], 2, 2) for l in labels]
|
||||
print("translated")
|
||||
for l in labels:
|
||||
print(l.bbox)
|
||||
|
||||
# print(labels)
|
||||
yield tile_reference, tile, labels, metadata
|
||||
|
||||
def split_respective_to_label(self, padding: int = 67):
|
||||
if self.labels is None:
|
||||
raise ValueError("No labels found. Only images having labels can be split.")
|
||||
|
||||
for i, label in enumerate(self.labels):
|
||||
tile_reference = f"_lbl-{i+1:02d}"
|
||||
# print(label.bbox)
|
||||
metadata = {"image_path": str(self.image_path), "label_path": str(self.label_path), "label_index": str(i)}
|
||||
|
||||
xc_norm, yc_norm, h_norm, w_norm = label.bbox # normalized coords
|
||||
xc, yc, h, w = [
|
||||
int(np.round(f))
|
||||
for f in [
|
||||
xc_norm * self.image.shape[1],
|
||||
yc_norm * self.image.shape[0],
|
||||
h_norm * self.image.shape[0],
|
||||
w_norm * self.image.shape[1],
|
||||
]
|
||||
] # image coords
|
||||
|
||||
# print("img coords:", xc, yc, h, w)
|
||||
pad_xneg = padding + 1 # int(w / 2) + padding
|
||||
pad_xpos = padding # int(w / 2) + padding
|
||||
pad_yneg = padding + 1 # int(h / 2) + padding
|
||||
pad_ypos = padding # int(h / 2) + padding
|
||||
if xc - pad_xneg < 0:
|
||||
pad_xneg = xc
|
||||
if pad_xpos + xc > self.image.shape[1]:
|
||||
pad_xpos = self.image.shape[1] - xc
|
||||
if yc - pad_yneg < 0:
|
||||
pad_yneg = yc
|
||||
if pad_ypos + yc > self.image.shape[0]:
|
||||
pad_ypos = self.image.shape[0] - yc
|
||||
|
||||
# print("pads:", pad_xneg, pad_xpos, pad_yneg, pad_ypos)
|
||||
|
||||
tile = self.image[
|
||||
yc - pad_yneg : yc + pad_ypos,
|
||||
xc - pad_xneg : xc + pad_xpos,
|
||||
]
|
||||
ny, nx = tile.shape
|
||||
x_offset = pad_xneg
|
||||
y_offset = pad_yneg
|
||||
|
||||
# print("tile shape:", tile.shape)
|
||||
|
||||
yolo_annotation = f"{label.class_id} " # {x_offset/nx} {y_offset/ny} {h_norm} {w_norm} "
|
||||
yolo_annotation += " ".join(
|
||||
[
|
||||
f"{(x*self.image.shape[1]-(xc - x_offset))/nx:.6f} {(y*self.image.shape[0]-(yc-y_offset))/ny:.6f}"
|
||||
for x, y in label.polygon
|
||||
]
|
||||
)
|
||||
print(yolo_annotation)
|
||||
new_label = Label(yolo_annotation=yolo_annotation)
|
||||
|
||||
yield tile_reference, tile, [new_label], metadata
|
||||
|
||||
|
||||
def main(args):
|
||||
|
||||
if args.output:
|
||||
args.output.mkdir(exist_ok=True, parents=True)
|
||||
(args.output / "images").mkdir(exist_ok=True)
|
||||
(args.output / "images-zoomed").mkdir(exist_ok=True)
|
||||
(args.output / "labels").mkdir(exist_ok=True)
|
||||
|
||||
for image_path in (args.input / "images").glob("*.tif"):
|
||||
data = ImageSplitter(
|
||||
image_path=image_path,
|
||||
label_path=(args.input / "labels" / image_path.stem).with_suffix(".txt"),
|
||||
)
|
||||
|
||||
if args.split_around_label:
|
||||
data = data.split_respective_to_label(padding=args.padding)
|
||||
else:
|
||||
data = data.split_into_tiles(patch_size=args.patch_size)
|
||||
|
||||
for tile_reference, tile, labels, metadata in data:
|
||||
print()
|
||||
print(tile_reference, tile.shape, labels, metadata) # len(labels) if labels else None)
|
||||
|
||||
# { debug
|
||||
debug = False
|
||||
if debug:
|
||||
plt.figure(figsize=(10, 10 * tile.shape[0] / tile.shape[1]))
|
||||
if labels is None:
|
||||
plt.imshow(tile, cmap="gray")
|
||||
plt.axis("off")
|
||||
plt.title(f"{image_path.name} ({tile_reference})")
|
||||
plt.show()
|
||||
continue
|
||||
|
||||
print(labels[0].bbox)
|
||||
# Draw annotations
|
||||
out = draw_annotations(
|
||||
cv2.cvtColor((tile / tile.max() * 255).astype(np.uint8), cv2.COLOR_GRAY2BGR),
|
||||
[l.to_string() for l in labels],
|
||||
alpha=0.1,
|
||||
)
|
||||
|
||||
# Convert BGR -> RGB for matplotlib display
|
||||
out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
|
||||
plt.imshow(out_rgb)
|
||||
plt.axis("off")
|
||||
plt.title(f"{image_path.name} ({tile_reference})")
|
||||
plt.show()
|
||||
# } debug
|
||||
|
||||
if args.output:
|
||||
# imwrite(args.output / "images" / f"{image_path.stem}_{tile_reference}.tif", tile, metadata=metadata)
|
||||
scale = 5
|
||||
tile_zoomed = zoom(tile, zoom=scale)
|
||||
metadata["scale"] = scale
|
||||
imwrite(
|
||||
args.output / "images" / f"{image_path.stem}_{tile_reference}.tif",
|
||||
tile_zoomed,
|
||||
metadata=metadata,
|
||||
imagej=True,
|
||||
)
|
||||
|
||||
if labels is not None:
|
||||
with open(args.output / "labels" / f"{image_path.stem}_{tile_reference}.txt", "w") as f:
|
||||
for label in labels:
|
||||
# label.offset_label(tile.shape[1], tile.shape[0])
|
||||
f.write(label.to_string() + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-i", "--input", type=Path)
|
||||
parser.add_argument("-o", "--output", type=Path)
|
||||
parser.add_argument(
|
||||
"-p",
|
||||
"--patch-size",
|
||||
nargs=2,
|
||||
type=int,
|
||||
default=[2, 2],
|
||||
help="Number of patches along height and width, rows and columns, respectively",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-sal",
|
||||
"--split-around-label",
|
||||
action="store_true",
|
||||
help="If enabled, the image will be split around the label and for each label, a separate image will be created.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--padding",
|
||||
type=int,
|
||||
default=67,
|
||||
help="Padding around the label when splitting around the label.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
main(args)
|
||||
1
src/utils/show_yolo_seg.py
Symbolic link
1
src/utils/show_yolo_seg.py
Symbolic link
@@ -0,0 +1 @@
|
||||
../../tests/show_yolo_seg.py
|
||||
157
src/utils/ultralytics_16bit_patch.py
Normal file
157
src/utils/ultralytics_16bit_patch.py
Normal file
@@ -0,0 +1,157 @@
|
||||
"""Ultralytics runtime patches for 16-bit TIFF training.
|
||||
|
||||
Goals:
|
||||
- Use `tifffile` to decode `.tif/.tiff` reliably (OpenCV can silently drop bit-depth depending on codec).
|
||||
- Preserve 16-bit data through the dataloader as `uint16` tensors.
|
||||
- Fix Ultralytics trainer normalization (default divides by 255) to scale `uint16` correctly.
|
||||
- Avoid uint8-forcing augmentations by recommending/setting hyp values (handled by caller).
|
||||
|
||||
This module is intended to be imported/called **before** instantiating/using YOLO.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from src.utils.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def apply_ultralytics_16bit_tiff_patches(*, force: bool = False) -> None:
|
||||
"""Apply runtime monkey-patches to Ultralytics to better support 16-bit TIFFs.
|
||||
|
||||
This function is safe to call multiple times.
|
||||
|
||||
Args:
|
||||
force: If True, re-apply patches even if already applied.
|
||||
"""
|
||||
|
||||
# Import inside function to ensure patching occurs before YOLO model/dataset is created.
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
# import tifffile
|
||||
import torch
|
||||
from src.utils.image import Image
|
||||
|
||||
from ultralytics.utils import patches as ul_patches
|
||||
|
||||
already_patched = getattr(ul_patches.imread, "__name__", "") == "tifffile_imread"
|
||||
if already_patched and not force:
|
||||
return
|
||||
|
||||
_original_imread = ul_patches.imread
|
||||
|
||||
def tifffile_imread(filename: str, flags: int = cv2.IMREAD_COLOR, pseudo_rgb: bool = True) -> Optional[np.ndarray]:
|
||||
"""Replacement for [`ultralytics.utils.patches.imread()`](venv/lib/python3.12/site-packages/ultralytics/utils/patches.py:20).
|
||||
|
||||
- For `.tif/.tiff`, uses `tifffile.imread()` and preserves dtype (e.g. uint16).
|
||||
- For other formats, falls back to Ultralytics' original implementation.
|
||||
- Always returns HWC (3 dims). For grayscale, returns (H, W, 1) or (H, W, 3) depending on requested flags.
|
||||
"""
|
||||
# print("here")
|
||||
# return _original_imread(filename, flags)
|
||||
ext = os.path.splitext(filename)[1].lower()
|
||||
if ext in (".tif", ".tiff"):
|
||||
arr = Image(filename).get_qt_rgb()[:, :, :3]
|
||||
|
||||
# Normalize common shapes:
|
||||
# - (H, W) -> (H, W, 1)
|
||||
# - (C, H, W) -> (H, W, C) (heuristic)
|
||||
if arr is None:
|
||||
return None
|
||||
if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[0] < arr.shape[1]:
|
||||
arr = np.transpose(arr, (1, 2, 0))
|
||||
if arr.ndim == 2:
|
||||
arr = arr[..., None]
|
||||
|
||||
# Ensure contiguous array for downstream OpenCV ops.
|
||||
# logger.info(f"Loading with monkey-patched imread: {filename}")
|
||||
arr = arr.astype(np.float32)
|
||||
arr /= arr.max()
|
||||
arr *= 2**8 - 1
|
||||
arr = arr.astype(np.uint8)
|
||||
# print(arr.shape, arr.dtype, any(np.isnan(arr).flatten()), np.where(np.isnan(arr)), arr.min(), arr.max())
|
||||
return np.ascontiguousarray(arr)
|
||||
|
||||
# logger.info(f"Loading with original imread: {filename}")
|
||||
return _original_imread(filename, flags)
|
||||
|
||||
# Patch the canonical reference.
|
||||
ul_patches.imread = tifffile_imread
|
||||
|
||||
# Patch common module-level imports (some Ultralytics modules do `from ... import imread`).
|
||||
# Importing these modules is safe and helps ensure the patched function is used.
|
||||
try:
|
||||
import ultralytics.data.base as _ul_base
|
||||
|
||||
_ul_base.imread = tifffile_imread
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
import ultralytics.data.loaders as _ul_loaders
|
||||
|
||||
_ul_loaders.imread = tifffile_imread
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Patch trainer normalization: default divides by 255 regardless of input dtype.
|
||||
from ultralytics.models.yolo.detect import train as detect_train
|
||||
|
||||
_orig_preprocess_batch = detect_train.DetectionTrainer.preprocess_batch
|
||||
|
||||
def preprocess_batch_16bit(self, batch: dict) -> dict: # type: ignore[override]
|
||||
# Start from upstream behavior to keep device placement + multiscale identical,
|
||||
# but replace the 255 division with dtype-aware scaling.
|
||||
# logger.info(f"Preprocessing batch with monkey-patched preprocess_batch")
|
||||
for k, v in batch.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
batch[k] = v.to(self.device, non_blocking=self.device.type == "cuda")
|
||||
|
||||
img = batch.get("img")
|
||||
if isinstance(img, torch.Tensor):
|
||||
# Decide scaling denom based on dtype (avoid expensive reductions if possible).
|
||||
if img.dtype == torch.uint8:
|
||||
denom = 255.0
|
||||
elif img.dtype == torch.uint16:
|
||||
denom = 65535.0
|
||||
elif img.dtype.is_floating_point:
|
||||
# Assume already in 0-1 range if float.
|
||||
denom = 1.0
|
||||
else:
|
||||
# Generic integer fallback.
|
||||
try:
|
||||
denom = float(torch.iinfo(img.dtype).max)
|
||||
except Exception:
|
||||
denom = 255.0
|
||||
|
||||
batch["img"] = img.float() / denom
|
||||
|
||||
# Multi-scale branch copied from upstream to avoid re-introducing `/255` scaling.
|
||||
if getattr(self.args, "multi_scale", False):
|
||||
import math
|
||||
import random
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
imgs = batch["img"]
|
||||
sz = (
|
||||
random.randrange(int(self.args.imgsz * 0.5), int(self.args.imgsz * 1.5 + self.stride))
|
||||
// self.stride
|
||||
* self.stride
|
||||
)
|
||||
sf = sz / max(imgs.shape[2:])
|
||||
if sf != 1:
|
||||
ns = [math.ceil(x * sf / self.stride) * self.stride for x in imgs.shape[2:]]
|
||||
imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
|
||||
batch["img"] = imgs
|
||||
|
||||
return batch
|
||||
|
||||
detect_train.DetectionTrainer.preprocess_batch = preprocess_batch_16bit
|
||||
|
||||
# Tag function to make it easier to detect patch state.
|
||||
setattr(detect_train.DetectionTrainer.preprocess_batch, "_ultralytics_16bit_patch", True)
|
||||
231
tests/show_yolo_seg.py
Normal file
231
tests/show_yolo_seg.py
Normal file
@@ -0,0 +1,231 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
show_yolo_seg.py
|
||||
|
||||
Usage:
|
||||
python show_yolo_seg.py /path/to/image.jpg /path/to/labels.txt
|
||||
|
||||
Supports:
|
||||
- Segmentation polygons: "class x1 y1 x2 y2 ... xn yn"
|
||||
- YOLO bbox lines as fallback: "class x_center y_center width height"
|
||||
Coordinates can be normalized [0..1] or absolute pixels (auto-detected).
|
||||
"""
|
||||
import sys
|
||||
import cv2
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
import random
|
||||
from shapely.geometry import LineString
|
||||
|
||||
from src.utils.image import Image
|
||||
|
||||
|
||||
def parse_label_line(line):
|
||||
parts = line.strip().split()
|
||||
if not parts:
|
||||
return None
|
||||
cls = int(float(parts[0]))
|
||||
coords = [float(x) for x in parts[1:]]
|
||||
return cls, coords
|
||||
|
||||
|
||||
def coords_are_normalized(coords):
|
||||
# If every coordinate is between 0 and 1 (inclusive-ish), assume normalized
|
||||
if not coords:
|
||||
return False
|
||||
return max(coords) <= 1.001
|
||||
|
||||
|
||||
def yolo_bbox_to_xyxy(coords, img_w, img_h):
|
||||
# coords: [xc, yc, w, h] normalized or absolute
|
||||
xc, yc, w, h = coords[:4]
|
||||
if max(coords) <= 1.001:
|
||||
xc *= img_w
|
||||
yc *= img_h
|
||||
w *= img_w
|
||||
h *= img_h
|
||||
x1 = int(round(xc - w / 2))
|
||||
y1 = int(round(yc - h / 2))
|
||||
x2 = int(round(xc + w / 2))
|
||||
y2 = int(round(yc + h / 2))
|
||||
return x1, y1, x2, y2
|
||||
|
||||
|
||||
def poly_to_pts(coords, img_w, img_h):
|
||||
# coords: [x1 y1 x2 y2 ...] either normalized or absolute
|
||||
if coords_are_normalized(coords[4:]):
|
||||
coords = [coords[i] * (img_w if i % 2 == 0 else img_h) for i in range(len(coords))]
|
||||
pts = np.array(coords, dtype=np.int32).reshape(-1, 2)
|
||||
return pts
|
||||
|
||||
|
||||
def random_color_for_class(cls):
|
||||
random.seed(cls) # deterministic per class
|
||||
return (
|
||||
0,
|
||||
0,
|
||||
255,
|
||||
) # tuple(int(x) for x in np.array([random.randint(0, 255) for _ in range(3)]))
|
||||
|
||||
|
||||
def draw_annotations(img, labels, alpha=0.4, draw_bbox_for_poly=True):
|
||||
# img: BGR numpy array
|
||||
overlay = img.copy()
|
||||
h, w = img.shape[:2]
|
||||
for line in labels:
|
||||
if isinstance(line, str):
|
||||
cls, coords = parse_label_line(line)
|
||||
if isinstance(line, tuple):
|
||||
cls, coords = line
|
||||
|
||||
if not coords:
|
||||
continue
|
||||
# polygon case (>=6 coordinates)
|
||||
if len(coords) >= 6:
|
||||
color = random_color_for_class(cls)
|
||||
|
||||
x1, y1, x2, y2 = yolo_bbox_to_xyxy(coords[:4], w, h)
|
||||
print(x1, y1, x2, y2)
|
||||
cv2.rectangle(img, (x1, y1), (x2, y2), color, 1)
|
||||
|
||||
pts = poly_to_pts(coords[4:], w, h)
|
||||
# line = LineString(pts)
|
||||
# # Buffer distance in pixels
|
||||
# buffered = line.buffer(3, cap_style=2, join_style=2)
|
||||
# coords = np.array(buffered.exterior.coords, dtype=np.int32)
|
||||
# cv2.fillPoly(overlay, [coords], color=(255, 255, 255))
|
||||
|
||||
# fill on overlay
|
||||
cv2.fillPoly(overlay, [pts], color)
|
||||
# outline on base image
|
||||
cv2.polylines(img, [pts], isClosed=True, color=color, thickness=1)
|
||||
# put class text at first point
|
||||
x, y = int(pts[0, 0]), int(pts[0, 1]) - 6
|
||||
if 0:
|
||||
cv2.putText(
|
||||
img,
|
||||
str(cls),
|
||||
(x, max(6, y)),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.6,
|
||||
(255, 255, 255),
|
||||
2,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
# YOLO bbox case (4 coords)
|
||||
elif len(coords) == 4:
|
||||
x1, y1, x2, y2 = yolo_bbox_to_xyxy(coords, w, h)
|
||||
color = random_color_for_class(cls)
|
||||
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
|
||||
cv2.putText(
|
||||
img,
|
||||
str(cls),
|
||||
(x1, max(6, y1 - 4)),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.6,
|
||||
(255, 255, 255),
|
||||
2,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
else:
|
||||
# Unknown / invalid format, skip
|
||||
continue
|
||||
|
||||
# blend overlay for filled polygons
|
||||
cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0, img)
|
||||
return img
|
||||
|
||||
|
||||
def load_labels_file(label_path):
|
||||
labels = []
|
||||
with open(label_path, "r") as f:
|
||||
for raw in f:
|
||||
line = raw.strip()
|
||||
if not line:
|
||||
continue
|
||||
parsed = parse_label_line(line)
|
||||
if parsed:
|
||||
labels.append(parsed)
|
||||
return labels
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Show YOLO segmentation / polygon annotations")
|
||||
parser.add_argument("image", type=str, help="Path to image file")
|
||||
parser.add_argument("--labels", type=str, help="Path to YOLO label file (polygons)")
|
||||
parser.add_argument("--alpha", type=float, default=0.4, help="Polygon fill alpha (0..1)")
|
||||
parser.add_argument("--no-bbox", action="store_true", help="Don't draw bounding boxes for polygons")
|
||||
args = parser.parse_args()
|
||||
|
||||
print(args)
|
||||
|
||||
img_path = Path(args.image)
|
||||
if args.labels:
|
||||
lbl_path = Path(args.labels)
|
||||
else:
|
||||
lbl_path = img_path.with_suffix(".txt")
|
||||
lbl_path = Path(str(lbl_path).replace("images", "labels"))
|
||||
|
||||
if not img_path.exists():
|
||||
print("Image not found:", img_path)
|
||||
sys.exit(1)
|
||||
if not lbl_path.exists():
|
||||
print("Label file not found:", lbl_path)
|
||||
sys.exit(1)
|
||||
|
||||
# img = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
|
||||
img = (Image(img_path).get_qt_rgb() * 255).astype(np.uint8)
|
||||
|
||||
if img is None:
|
||||
print("Could not load image:", img_path)
|
||||
sys.exit(1)
|
||||
|
||||
labels = load_labels_file(str(lbl_path))
|
||||
if not labels:
|
||||
print("No labels parsed from", lbl_path)
|
||||
# continue and just show image
|
||||
out = draw_annotations(img.copy(), labels, alpha=args.alpha, draw_bbox_for_poly=(not args.no_bbox))
|
||||
|
||||
out_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
plt.figure(figsize=(10, 10 * out.shape[0] / out.shape[1]))
|
||||
if 0:
|
||||
plt.imshow(out_rgb.transpose(1, 0, 2))
|
||||
else:
|
||||
plt.imshow(out_rgb)
|
||||
|
||||
for label in labels:
|
||||
lclass, coords = label
|
||||
# print(lclass, coords)
|
||||
bbox = coords[:4]
|
||||
# print("bbox", bbox)
|
||||
bbox = np.array(bbox) * np.array([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
|
||||
yc, xc, h, w = bbox
|
||||
# print("bbox", bbox)
|
||||
|
||||
# polyline = np.array(coords[4:]).reshape(-1, 2) * np.array([img.shape[1], img.shape[0]])
|
||||
polyline = np.array(coords).reshape(-1, 2) * np.array([img.shape[1], img.shape[0]])
|
||||
# print("pl", coords[4:])
|
||||
# print("pl", polyline)
|
||||
|
||||
# Convert BGR -> RGB for matplotlib display
|
||||
# out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
|
||||
# out_rgb = Image()
|
||||
plt.plot(polyline[:, 0], polyline[:, 1], "y", linewidth=2)
|
||||
if 0:
|
||||
plt.plot(
|
||||
[yc - h / 2, yc - h / 2, yc + h / 2, yc + h / 2, yc - h / 2],
|
||||
[xc - w / 2, xc + w / 2, xc + w / 2, xc - w / 2, xc - w / 2],
|
||||
"r",
|
||||
linewidth=2,
|
||||
)
|
||||
|
||||
# plt.axis("off")
|
||||
plt.title(f"{img_path.name} ({lbl_path.name})")
|
||||
plt.show()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user