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)