ScanSlicer/scanslicer/controller_image_io.py

import ntpath
import os

import cv2
import numpy as np
from PySide6.QtWidgets import QFileDialog, QMessageBox


class ScanSlicerImageIOMixin:
    def resolve_directory_dialog_start_dir(self):
        if self.input_dir and os.path.isdir(self.input_dir):
            return self.input_dir

        # In containers, default to mounted host paths when available.
        if os.path.exists("/.dockerenv"):
            for candidate in ("/host-home", "/work"):
                if os.path.isdir(candidate):
                    return candidate

        home_dir = os.path.expanduser("~")
        if os.path.isdir(home_dir):
            return home_dir
        return ""

    def refresh_goto_image_dropdown(self):
        dropdown = getattr(self, "goto_image_dropdown", None)
        if dropdown is None:
            return

        dropdown.blockSignals(True)
        dropdown.clear()
        for image_path in self.source_images:
            dropdown.addItem(ntpath.basename(image_path), image_path)
        dropdown.setEnabled(len(self.source_images) > 0)
        dropdown.blockSignals(False)
        self.sync_goto_image_selection()

    def sync_goto_image_selection(self):
        dropdown = getattr(self, "goto_image_dropdown", None)
        if dropdown is None:
            return

        has_images = len(self.source_images) > 0
        dropdown.setEnabled(has_images)
        dropdown.blockSignals(True)
        if not has_images:
            dropdown.setCurrentIndex(-1)
            dropdown.blockSignals(False)
            return

        index = 0 if self.source_index is None else int(self.source_index)
        index = max(0, min(index, len(self.source_images) - 1))
        dropdown.setCurrentIndex(index)
        dropdown.blockSignals(False)

    def goto_image_from_dropdown(self, index):
        if self.processing_busy:
            return
        if index < 0 or index >= len(self.source_images):
            return
        if self.source_index == index:
            return

        self.source_index = index
        self.show_current_image(new_image=True)

    def sort_image_paths(self, image_paths):
        def safe_mtime(path):
            try:
                return os.path.getmtime(path)
            except OSError:
                return float("inf")

        if self.sort_mode == "time":
            return sorted(
                image_paths,
                key=lambda path: (
                    safe_mtime(path),
                    ntpath.basename(path).lower(),
                    path.lower(),
                ),
            )
        return sorted(
            image_paths,
            key=lambda path: (
                ntpath.basename(path).lower(),
                path.lower(),
            ),
        )

    def show_current_image(self, new_image=True):
        if self.processing_busy:
            return
        if self.source_index is None or len(self.source_images) == 0:
            self.sync_goto_image_selection()
            return

        self.autoslicer.load_image(self.source_images[self.source_index])
        if self.autoslicer.image_loaded():
            self.update_status_filename(self.source_images[self.source_index])
            self.sync_goto_image_selection()
            self.invalidate_preview_cache()
            self.update_preview(new_image=new_image, run_detection=True)
            self.update_statusbar("Loaded " + self.source_images[self.source_index])

    def load_image(self, move_index=0):
        if self.processing_busy:
            return

        if len(self.source_images) == 0:
            self.open_directory()
            if len(self.source_images) == 0:
                return

        if self.source_index is None:
            self.source_index = 0
        else:
            self.source_index += move_index

            if self.source_index < 0:
                self.source_index = 0
                QMessageBox.warning(self, "No previous", "This is the first image")
                return

            if self.source_index >= len(self.source_images):
                self.source_index = len(self.source_images) - 1
                QMessageBox.warning(self, "No next", "This is the last image")
                return

        self.show_current_image(new_image=True)

    def set_save_format(self, choice):
        self.save_format = choice
        self.update_quality_controls_state()

    def set_jpg_quality(self, value):
        self.jpg_quality = int(value)

    def set_png_compression(self, value):
        self.png_compression = int(value)

    def update_quality_controls_state(self):
        jpeg_enabled = self.save_format == "jpg"
        png_enabled = self.save_format == "png"
        self.jpg_quality_label.setVisible(jpeg_enabled)
        self.jpg_quality_spinbox.setEnabled(jpeg_enabled)
        self.jpg_quality_spinbox.setVisible(jpeg_enabled)
        self.png_compression_label.setVisible(png_enabled)
        self.png_compression_spinbox.setEnabled(png_enabled)
        self.png_compression_spinbox.setVisible(png_enabled)

    def set_sort_mode(self, choice):
        mode = choice.lower()
        if mode not in ("name", "time"):
            return
        if mode == self.sort_mode:
            return

        self.sort_mode = mode
        if len(self.source_images) == 0:
            return

        current_image = None
        if self.source_index is not None and 0 <= self.source_index < len(self.source_images):
            current_image = self.source_images[self.source_index]

        self.source_images = self.sort_image_paths(self.source_images)
        self.refresh_goto_image_dropdown()
        if current_image is not None and current_image in self.source_images:
            self.source_index = self.source_images.index(current_image)
        else:
            self.source_index = 0

        self.show_current_image(new_image=True)

    def update_output_dir_label(self):
        if self.output_dir:
            self.output_dir_label.setText(self.output_dir)
        else:
            self.output_dir_label.setText("(not set)")

    def set_output_directory(self, basedir=None):
        # QPushButton.clicked emits a bool; ignore it when this method is used as a slot.
        if isinstance(basedir, bool):
            basedir = None

        if basedir is None:
            start_dir = self.output_dir if self.output_dir else (self.input_dir if self.input_dir else "")
            basedir = QFileDialog.getExistingDirectory(self, "Select output directory", start_dir)

        if basedir is None or len(basedir) == 0:
            return

        self.output_dir = basedir
        self.update_output_dir_label()
        self.update_statusbar("Output directory set to " + self.output_dir)

    def save_all(self):
        if self.processing_busy:
            return

        if self.source_images is None or self.source_index is None or self.source_index < 0:
            QMessageBox.warning(self, "No image loaded", "Load an image first")
            return

        basedir = self.output_dir
        if basedir is None or len(basedir) == 0:
            QMessageBox.warning(self, "No output directory", "Set an output directory first")
            return

        if not os.path.isdir(basedir):
            QMessageBox.warning(self, "Invalid output directory", f"Directory does not exist:\n{basedir}")
            return

        total_saved = 0
        saved_boxes = []
        save_params = None
        if self.save_format == "jpg":
            save_params = [int(cv2.IMWRITE_JPEG_QUALITY), self.jpg_quality]
        elif self.save_format == "png":
            save_params = [int(cv2.IMWRITE_PNG_COMPRESSION), self.png_compression]

        self.set_processing_busy(True, "Saving images...")
        try:
            for i, slice_data in enumerate(self.slicing_canvas.slices):
                if not slice_data.locked:
                    continue

                basename = ntpath.basename(self.source_images[self.source_index])
                basename = os.path.splitext(basename)[0] + '_' + f'{i:03}' + '.' + self.save_format

                outname = basedir + os.path.sep + basename
                self.autoslicer.save_slice(slice_data.bbox, outname, save_params)
                self.update_statusbar("Saved " + outname)
                total_saved += 1
                saved_boxes.append(np.array(slice_data.bbox, dtype=np.float64).reshape(4, 2))
        finally:
            self.set_processing_busy(False)

        if total_saved == 0:
            QMessageBox.warning(
                self,
                "No locked slice to save!",
                "Select a slice by clicking inside it, then lock it using the toolbar button.",
            )
        else:
            try:
                image_path = self.source_images[self.source_index]
                feedback_saved = self.autoslicer.record_feedback(
                    image_path=image_path,
                    final_boxes=saved_boxes,
                    detected_boxes=self.last_detected_boxes,
                )
                if feedback_saved:
                    self.append_status_log(
                        f"Feedback recorded: image={image_path}, final_boxes={len(saved_boxes)}"
                    )
            except OSError as exc:
                self.append_status_log(f"Feedback save failed: {exc}")
            QMessageBox.information(self, "Slices saved", f"{total_saved} slices have been saved")

    def open_directory(self, basedir=None):
        if self.processing_busy:
            return
        # QPushButton.clicked emits a bool; ignore it when this method is used as a slot.
        if isinstance(basedir, bool):
            basedir = None

        if basedir is None:
            start_dir = self.resolve_directory_dialog_start_dir()
            basedir = QFileDialog.getExistingDirectory(self, "Open directory", start_dir)

        if basedir is None or len(basedir) == 0:
            return

        extensions = (".png", ".jpg", ".jpeg")
        discovered_images = []
        with os.scandir(basedir) as entries:
            for entry in entries:
                if not entry.is_file():
                    continue
                if entry.name.lower().endswith(extensions):
                    discovered_images.append(entry.path)

        self.source_images = self.sort_image_paths(discovered_images)
        self.refresh_goto_image_dropdown()
        self.input_dir = basedir
        self.output_dir = basedir
        self.update_output_dir_label()
        self.source_index = None

        if len(self.source_images) == 0:
            self.update_status_filename(None)
            self.update_statusbar("No images found in selected directory.")
        else:
            self.load_image(0)

    def next_image(self):
        self.load_image(1)

    def prev_image(self):
        self.load_image(-1)