imagebaker 0.0.49__py3-none-any.whl → 0.0.51__py3-none-any.whl

imagebaker/list_views/layer_settings.py

@@ -20,7 +20,14 @@ class LayerSettings(QDockWidget):
     layerState = Signal(LayerState)
     messageSignal = Signal(str)
 
-    def __init__(self, parent=None, max_xpos=1000, max_ypos=1000, max_scale=100):
+    def __init__(
+        self,
+        parent=None,
+        max_xpos=1000,
+        max_ypos=1000,
+        max_scale=100,
+        max_edge_width=10,
+    ):
         super().__init__("BaseLayer Settings", parent)
         self.selected_layer: BaseLayer = None
 
@@ -29,6 +36,7 @@ class LayerSettings(QDockWidget):
         self.max_xpos = max_xpos
         self.max_ypos = max_ypos
         self.max_scale = max_scale
+        self.max_edge_width = max_edge_width
         self.init_ui()
         self.setFeatures(
             QDockWidget.DockWidgetMovable | QDockWidget.DockWidgetFloatable
@@ -67,6 +75,12 @@ class LayerSettings(QDockWidget):
         self.main_layout.addWidget(self.scale_y_slider["widget"])
         self.rotation_slider = self.create_slider("Rotation:", 0, 360, 0, 1)
         self.main_layout.addWidget(self.rotation_slider["widget"])
+        self.edge_opacity_slider = self.create_slider("Edge Opacity:", 0, 255, 255, 1)
+        self.main_layout.addWidget(self.edge_opacity_slider["widget"])
+        self.edge_width_slider = self.create_slider(
+            "Edge Width:", 0, self.max_edge_width, 5, 1
+        )
+        self.main_layout.addWidget(self.edge_width_slider["widget"])
 
         # Add stretch to push content to the top
         self.main_layout.addStretch()
@@ -125,7 +139,6 @@ class LayerSettings(QDockWidget):
 
         try:
             self._disable_updates = True
-
             if sender == self.opacity_slider["slider"]:
                 self.selected_layer.opacity = value
             elif sender == self.x_slider["slider"]:
@@ -138,6 +151,12 @@ class LayerSettings(QDockWidget):
                 self.selected_layer.scale_y = value / 100.0
             elif sender == self.rotation_slider["slider"]:
                 self.selected_layer.rotation = value
+            elif sender == self.edge_opacity_slider["slider"]:
+                self.selected_layer.edge_opacity = value
+                self.selected_layer._apply_edge_opacity()
+            elif sender == self.edge_width_slider["slider"]:
+                self.selected_layer.edge_width = value
+                self.selected_layer._apply_edge_opacity()
 
             self.selected_layer.update()  # Trigger a repaint
 
@@ -154,6 +173,10 @@ class LayerSettings(QDockWidget):
             rotation=self.selected_layer.rotation,
             scale_x=self.selected_layer.scale_x,
             scale_y=self.selected_layer.scale_y,
+            opacity=self.selected_layer.opacity,
+            edge_opacity=self.selected_layer.edge_opacity,
+            edge_width=self.selected_layer.edge_width,
+            visible=self.selected_layer.visible,
         )
         logger.info(f"Storing state {bake_settings}")
         self.messageSignal.emit(f"Stored state for {bake_settings.layer_name}")
@@ -211,6 +234,12 @@ class LayerSettings(QDockWidget):
                 self.rotation_slider["slider"].setValue(
                     int(self.selected_layer.rotation)
                 )
+                self.edge_opacity_slider["slider"].setValue(
+                    int(self.selected_layer.edge_opacity)
+                )
+                self.edge_width_slider["slider"].setValue(
+                    int(self.selected_layer.edge_width)
+                )
             else:
                 self.widget.setEnabled(False)
                 self.layer_name_label.setText("No BaseLayer")

imagebaker/models/base_model.py

@@ -39,7 +39,7 @@ class BaseModel(ABC):
         return image
 
     # @abstractmethod
-    def postprocess(self, output) -> PredictionResult:
+    def postprocess(self, output) -> list[PredictionResult]:
         return output
 
     def predict(

imagebaker/tabs/baker_tab.py

@@ -66,6 +66,7 @@ class BakerTab(QWidget):
             max_xpos=self.config.max_xpos,
             max_ypos=self.config.max_ypos,
             max_scale=self.config.max_scale,
+            max_edge_width=self.config.max_edge_width,
         )
         self.layer_list = LayerList(
             canvas=self.current_canvas,
@@ -415,15 +416,7 @@ class BakerTab(QWidget):
         """Seek to a specific state using the timeline slider."""
         self.messageSignal.emit(f"Seeking to step {step}")
         logger.info(f"Seeking to step {step}")
-
-        # Get the states for the selected step
-        if step in self.current_canvas.states:
-            states = self.current_canvas.states[step]
-            for state in states:
-                layer = self.current_canvas.get_layer(state.layer_id)
-                if layer:
-                    layer.layer_state = state
-                    layer.update()
+        self.current_canvas.seek_state(step)
 
         # Update the canvas
         self.current_canvas.update()

imagebaker/tabs/layerify_tab.py

@@ -40,6 +40,7 @@ from imagebaker.core.defs import (
 from collections import deque
 from typing import Deque
 from dataclasses import dataclass
+import os
 
 
 @dataclass
@@ -103,13 +104,17 @@ class LayerifyTab(QWidget):
         """Connect all necessary signals"""
         # Connect all layers in the deque to annotation list
         for layer in self.annotable_layers:
-            layer.annotationAdded.connect(self.annotation_list.update_list)
-            layer.annotationUpdated.connect(self.annotation_list.update_list)
+            # layer.annotationAdded.connect(self.annotation_list.update_list)
+            layer.annotationAdded.connect(self.on_annotation_added)
+            # layer.annotationUpdated.connect(self.annotation_list.update_list)
+            layer.annotationUpdated.connect(self.on_annotation_updated)
             layer.messageSignal.connect(self.messageSignal)
             layer.layerSignal.connect(self.add_layer)
+            layer.labelUpdated.connect(self.on_label_update)
 
         # Connect image list panel signals
         self.image_list_panel.imageSelected.connect(self.on_image_selected)
+        self.image_list_panel.activeImageEntries.connect(self.update_active_entries)
 
     def init_ui(self):
         """Initialize the UI components"""
@@ -118,7 +123,9 @@ class LayerifyTab(QWidget):
             None, parent=self.main_window, max_name_length=self.config.max_name_length
         )
         self.image_list_panel = ImageListPanel(
-            self.image_entries, self.processed_images
+            self.image_entries,
+            self.processed_images,
+            images_per_page=self.config.deque_maxlen,
         )
 
         self.main_window.addDockWidget(Qt.LeftDockWidgetArea, self.image_list_panel)
@@ -166,6 +173,8 @@ class LayerifyTab(QWidget):
         if not image_entry.is_baked_result:  # Regular image
             image_path = image_entry.data
             self.curr_image_idx = self.image_entries.index(image_entry)
+            # convert curr_image_idx to correct index
+            self.curr_image_idx = self.curr_image_idx % len(self.annotable_layers)
 
             # Make the corresponding layer visible and set the image
             selected_layer = self.annotable_layers[self.curr_image_idx]
@@ -206,6 +215,8 @@ class LayerifyTab(QWidget):
             for i, layer in enumerate(self.annotable_layers):
                 if i < len(self.image_entries):
                     layer.set_image(self.image_entries[i].data)
+                    self.load_layer_annotations(layer)
+                    layer.layer_name = f"Layer_{i + 1}"
                     layer.setVisible(
                         i == 0
                     )  # Only the first layer is visible by default
@@ -224,12 +235,68 @@ class LayerifyTab(QWidget):
         self.image_list_panel.update_image_list(self.image_entries)
         self.update()
 
+    def save_layer_annotations(self, layer: AnnotableLayer):
+        """Save annotations for a specific layer"""
+        if len(layer.annotations) > 0:
+            file_path = layer.file_path
+            file_name = file_path.name
+            save_dir = self.config.cache_dir / f"{file_name}.json"
+            Annotation.save_as_json(layer.annotations, save_dir)
+            logger.info(f"Saved annotations for {layer.layer_name} to {save_dir}")
+
+    def load_layer_annotations(self, layer: AnnotableLayer):
+        """Load annotations for a specific layer"""
+        if layer.file_path:
+            file_path = layer.file_path
+            file_name = file_path.name
+            load_dir = self.config.cache_dir / f"{file_name}.json"
+            if load_dir.exists():
+                layer.annotations = Annotation.load_from_json(load_dir)
+                logger.info(
+                    f"Loaded annotations for {layer.layer_name} from {load_dir}"
+                )
+            else:
+                logger.warning(f"No annotations found for {layer.layer_name}")
+
+    def update_active_entries(self, image_entries: list[ImageEntry]):
+        """Update the active entries in the image list panel."""
+        self.curr_image_idx = 0
+        for i, layer in enumerate(self.annotable_layers):
+            self.save_layer_annotations(layer)
+            layer.annotations = []
+
+            if i < len(image_entries):
+                # get index on the self.image_entries
+                idx = self.image_entries.index(image_entries[i])
+                if self.image_entries[idx].is_baked_result:
+                    # if the image is a baked result, set the layer to the baked result
+                    layer = self.image_entries[idx].data
+                    layer.file_path = layer.file_path
+                else:
+                    layer.set_image(self.image_entries[idx].data)
+                self.load_layer_annotations(layer)
+
+                layer.layer_name = f"Layer_{idx + 1}"
+                layer.setVisible(i == 0)
+                if i == 0:
+                    self.layer = layer
+            else:
+                layer.setVisible(False)
+        logger.info("Updated active entries in image list panel.")
+
     def clear_annotations(self):
         """Safely clear all annotations"""
         try:
             # Clear layer annotations
             self.clearAnnotations.emit()
             self.messageSignal.emit("Annotations cleared")
+            # clear cache annotation of layer
+            annotation_path = (
+                self.config.cache_dir / f"{self.layer.file_path.name}.json"
+            )
+            if annotation_path.exists():
+                os.remove(annotation_path)
+                logger.info(f"Cleared annotations from {annotation_path}")
 
         except Exception as e:
             logger.error(f"Clear error: {str(e)}")
@@ -241,13 +308,17 @@ class LayerifyTab(QWidget):
         Args:
             annotation (Annotation): The annotation that was added.
         """
-        if annotation.label not in self.config.predefined_labels:
+
+        # if annotation.label is not in the predefined labels, add it
+        if annotation.label not in [lbl.name for lbl in self.config.predefined_labels]:
+            logger.info(f"Label {annotation.label} created.")
             self.config.predefined_labels.append(
                 Label(annotation.label, annotation.color)
             )
             self.update_label_combo()
         logger.info(f"Added annotation: {annotation.label}")
         self.messageSignal.emit(f"Added annotation: {annotation.label}")
+        self.save_layer_annotations(self.layer)
 
         # Refresh the annotation list
         self.annotation_list.update_list()
@@ -259,11 +330,12 @@ class LayerifyTab(QWidget):
         Args:
             annotation (Annotation): The updated annotation.
         """
-        logger.info(f"Updated annotation: {annotation.label}")
+        # logger.info(f"Updated annotation: {annotation}")
         self.messageSignal.emit(f"Updated annotation: {annotation.label}")
 
         # Refresh the annotation list
         self.annotation_list.update_list()
+        self.save_layer_annotations(self.layer)
 
     def update_label_combo(self):
         """
@@ -276,6 +348,27 @@ class LayerifyTab(QWidget):
             pixmap = QPixmap(16, 16)
             pixmap.fill(label.color)
             self.label_combo.addItem(QIcon(pixmap), label.name)
+        logger.info("Updated label combo box with predefined labels.")
+        self.label_combo.setCurrentText(self.current_label)
+
+    def on_label_update(self, old_new_label: tuple[str, str]):
+        new_labels = []
+        index = 0
+        for i, label in enumerate(self.config.predefined_labels):
+            if label.name == old_new_label[0]:
+                label.name = old_new_label[1]
+                index = i
+            new_labels.append(label)
+
+        self.config.predefined_labels = new_labels
+        logger.info(f"Updated label from {old_new_label[0]} to {old_new_label[1]}")
+        self.messageSignal.emit(
+            f"Updated label from {old_new_label[0]} to {old_new_label[1]}."
+        )
+
+        self.update_label_combo()
+        self.handle_label_change(index=index)
+        self.label_combo.update()
 
     def load_default_image(self):
         """
@@ -463,6 +556,7 @@ class LayerifyTab(QWidget):
     def handle_model_error(self, error):
         logger.error(f"Model error: {error}")
         QMessageBox.critical(self, "Error", f"Model error: {error}")
+        self.loading_dialog.close()
 
     def save_annotations(self):
         """Save annotations to a JSON file."""
@@ -588,6 +682,24 @@ class LayerifyTab(QWidget):
         )
         msg = f"Label changed to {self.current_label}"
         self.messageSignal.emit(msg)
+        self.layer.selected_annotation = self.layer._get_selected_annotation()
+        if self.layer.selected_annotation:
+            annotations = []
+            for ann in self.layer.annotations:
+                if ann == self.layer.selected_annotation:
+                    ann.label = label_info.name
+                    ann.color = label_info.color
+                annotations.append(ann)
+
+            self.layer.annotations = annotations
+            self.on_annotation_updated(self.layer.selected_annotation)
+            # disable label change callback
+            self.label_combo.currentIndexChanged.disconnect()
+            self.label_combo.currentIndexChanged.connect(lambda: None)
+            self.update_label_combo()
+            # set it back
+            self.label_combo.currentIndexChanged.connect(self.handle_label_change)
+
         self.layer.update()
         self.update()
 
@@ -627,7 +739,7 @@ class LayerifyTab(QWidget):
             ("🎨", "Color", self.choose_color),
             ("🧅", "Layerify All", self.layerify_all),
             ("🏷️", "Add Label", self.add_new_label),
-            ("🗑️", "Clear", lambda x: self.clearAnnotations.emit()),
+            ("🗑️", "Clear", self.clear_annotations),
         ]
 
         # Folder navigation buttons
@@ -635,18 +747,6 @@ class LayerifyTab(QWidget):
         self.select_folder_btn.clicked.connect(self.select_folder)
         toolbar_layout.addWidget(self.select_folder_btn)
 
-        self.next_image_btn = QPushButton("Next")
-        self.next_image_btn.clicked.connect(self.show_next_image)
-        toolbar_layout.addWidget(self.next_image_btn)
-
-        self.prev_image_btn = QPushButton("Prev")
-        self.prev_image_btn.clicked.connect(self.show_prev_image)
-        toolbar_layout.addWidget(self.prev_image_btn)
-
-        # Initially hide next/prev buttons
-        self.next_image_btn.setVisible(False)
-        self.prev_image_btn.setVisible(False)
-
         # Add mode buttons
         for icon, text, mode in modes:
             btn_txt = icon + text
@@ -702,11 +802,25 @@ class LayerifyTab(QWidget):
                     ImageEntry(is_baked_result=False, data=img_path)
                 )
 
+        # load from bake folder if it exists
+        bake_folder = self.config.bake_dir
+        if bake_folder.exists() and bake_folder.is_dir():
+            for img_path in bake_folder.glob("*.*"):
+                if img_path.suffix.lower() in [
+                    ".jpg",
+                    ".jpeg",
+                    ".png",
+                    ".bmp",
+                    ".tiff",
+                ]:
+                    self.image_entries.append(
+                        ImageEntry(is_baked_result=False, data=img_path)
+                    )
+
     def select_folder(self):
         """Allow the user to select a folder and load images from it."""
         folder_path = QFileDialog.getExistingDirectory(self, "Select Folder")
         if folder_path:
-            self.image_entries = []  # Clear the existing image paths
             folder_path = Path(folder_path)
 
             self._load_images_from_folder(folder_path)
@@ -725,9 +839,6 @@ class LayerifyTab(QWidget):
                 # Load the first set of images into the layers
                 self.load_default_images()
 
-                # Unhide the next/prev buttons if there are multiple images
-                self.next_image_btn.setVisible(len(self.image_entries) > 1)
-                self.prev_image_btn.setVisible(len(self.image_entries) > 1)
             else:
                 QMessageBox.warning(
                     self,
@@ -735,28 +846,6 @@ class LayerifyTab(QWidget):
                     "No valid image files found in the selected folder.",
                 )
 
-    def show_next_image(self):
-        """Show next image in the list. If at the end, show first image."""
-        if self.curr_image_idx < len(self.image_entries) - 1:
-            self.curr_image_idx += 1
-        else:
-            self.curr_image_idx = 0
-        self.layer.set_image(self.image_entries[self.curr_image_idx]["data"])
-        self.messageSignal.emit(
-            f"Showing image {self.curr_image_idx + 1}/{len(self.image_entries)}"
-        )
-
-    def show_prev_image(self):
-        """Show previous image in the list. If at the start, show last image."""
-        if self.curr_image_idx > 0:
-            self.curr_image_idx -= 1
-        else:
-            self.curr_image_idx = len(self.image_entries) - 1
-        self.layer.set_image(self.image_entries[self.curr_image_idx]["data"])
-        self.messageSignal.emit(
-            f"Showing image {self.curr_image_idx + 1}/{len(self.image_entries)}"
-        )
-
     def __del__(self):
         logger.warning(f"Tab {id(self)} deleted")
 
@@ -769,10 +858,20 @@ class LayerifyTab(QWidget):
             config=self.config,
             canvas_config=self.canvas_config,
         )
-        layer.annotations = baking_result.annotations
+        # save it in cache
+        filename = baking_result.filename
+        filepath = self.config.bake_dir / filename.name
+        baking_result.image.save(str(filepath))
+        #
+        Annotation.save_as_json(
+            baking_result.annotations, self.config.cache_dir / f"{filename.name}.json"
+        )
 
-        layer.annotationAdded.connect(self.annotation_list.update_list)
-        layer.annotationUpdated.connect(self.annotation_list.update_list)
+        layer.set_image(filepath)
+
+        layer.annotationAdded.connect(self.on_annotation_added)
+        layer.annotationUpdated.connect(self.on_annotation_updated)
+        layer.labelUpdated.connect(self.on_label_update)
         layer.messageSignal.connect(self.messageSignal)
         layer.layerSignal.connect(self.add_layer)
 
@@ -784,7 +883,7 @@ class LayerifyTab(QWidget):
         self.annotable_layers.append(layer)
 
         # Add baked result to image_entries
-        baked_result_entry = ImageEntry(is_baked_result=True, data=layer)
+        baked_result_entry = ImageEntry(is_baked_result=False, data=filepath)
         self.image_entries.append(baked_result_entry)
         # baking_result.image.save(str(baking_result.filename))
         layer.update()
@@ -792,6 +891,12 @@ class LayerifyTab(QWidget):
         logger.info("A baked result has arrived, adding it to the image list.")
 
         # Update the image list panel
+        # find the page index where this layer is
+        page_index = (
+            self.image_entries.index(baked_result_entry) // self.config.deque_maxlen
+        )
+        # set the current page to the page index
+        self.image_list_panel.current_page = page_index
         self.image_list_panel.update_image_list(self.image_entries)
         self.image_list_panel.imageSelected.emit(baked_result_entry)
 
@@ -830,6 +935,19 @@ class LayerifyTab(QWidget):
                 self.annotation_list.update_list()
                 logger.info("Selected annotation deleted.")
 
+        # if clicked q, set the mode to point
+        elif key == Qt.Key_Q:
+            self.layer.set_mode(MouseMode.POINT)
+            logger.info("Mouse mode set to POINT.")
+        # if clicked w, set the mode to polygon
+        elif key == Qt.Key_W:
+            self.layer.set_mode(MouseMode.POLYGON)
+            logger.info("Mouse mode set to POLYGON.")
+        # if clicked e, set the mode to rectangle
+        elif key == Qt.Key_E:
+            self.layer.set_mode(MouseMode.RECTANGLE)
+            logger.info("Mouse mode set to RECTANGLE.")
+
         # Pass the event to the annotation list if it needs to handle it
         if self.annotation_list.hasFocus():
             self.annotation_list.keyPressEvent(event)

imagebaker/utils/__init__.py

@@ -0,0 +1,3 @@
+from .utils import *  # noqa: F401, F403
+from .transform_mask import *  # noqa: F401, F403
+from .vis import *  # noqa: F401, F403

imagebaker/utils/image.py

@@ -49,8 +49,11 @@ def draw_annotations(image: np.ndarray, annotations: list[Annotation]) -> np.nda
     Returns:
         np.ndarray: Image with annotations drawn.
     """
+    color = (0, 255, 0, 255) if image.shape[2] == 4 else (0, 255, 0)
+
     for i, ann in enumerate(annotations):
         if ann.rectangle:
+            # if image has alpha channel, make color full alpha
             cv2.rectangle(
                 image,
                 (int(ann.rectangle.x()), int(ann.rectangle.y())),
@@ -58,7 +61,7 @@ def draw_annotations(image: np.ndarray, annotations: list[Annotation]) -> np.nda
                     int(ann.rectangle.x() + ann.rectangle.width()),
                     int(ann.rectangle.y() + ann.rectangle.height()),
                 ),
-                (0, 255, 0),
+                color,
                 2,
             )
             rect_center = ann.rectangle.center()
@@ -69,7 +72,7 @@ def draw_annotations(image: np.ndarray, annotations: list[Annotation]) -> np.nda
                 (int(rect_center.x()), int(rect_center.y())),
                 cv2.FONT_HERSHEY_SIMPLEX,
                 1,
-                (0, 255, 0),
+                color,
                 2,
             )
         elif ann.polygon:
@@ -77,7 +80,7 @@ def draw_annotations(image: np.ndarray, annotations: list[Annotation]) -> np.nda
                 image,
                 [np.array([[int(p.x()), int(p.y())] for p in ann.polygon])],
                 True,
-                (0, 255, 0),
+                color,
                 2,
             )
             polygon_center = ann.polygon.boundingRect().center()
@@ -87,19 +90,19 @@ def draw_annotations(image: np.ndarray, annotations: list[Annotation]) -> np.nda
                 (int(polygon_center.x()), int(polygon_center.y())),
                 cv2.FONT_HERSHEY_SIMPLEX,
                 1,
-                (0, 255, 0),
+                color,
                 2,
             )
         elif ann.points:
             for p in ann.points:
-                cv2.circle(image, (int(p.x()), int(p.y())), 5, (0, 255, 0), -1)
+                cv2.circle(image, (int(p.x()), int(p.y())), 5, color, -1)
             cv2.putText(
                 image,
                 ann.label,
                 (int(ann.points[0].x()), int(ann.points[0].y())),
                 cv2.FONT_HERSHEY_SIMPLEX,
                 1,
-                (0, 255, 0),
+                color,
                 2,
             )
     return image

imagebaker/utils/state_utils.py

@@ -59,9 +59,13 @@ def calculate_intermediate_states(
             visible=current_state.visible,
             allow_annotation_export=current_state.allow_annotation_export,
             playing=current_state.playing,
-            selected=current_state.selected,
+            selected=False,
             is_annotable=current_state.is_annotable,
             status=current_state.status,
+            edge_opacity=previous_state.edge_opacity
+            + (current_state.edge_opacity - previous_state.edge_opacity) * (i / steps),
+            edge_width=previous_state.edge_width
+            + (current_state.edge_width - previous_state.edge_width) * (i / steps),
         )
 
         # Deep copy the drawing_states from the previous_state

imagebaker/utils/utils.py

@@ -0,0 +1,26 @@
+import numpy as np
+import cv2
+
+
+def generate_color_map(num_colors: int = 20):
+    """Generate a color map for the segmentation masks"""
+    np.random.seed(42)  # For reproducible colors
+
+    colors = {}
+    for i in range(num_colors):
+        # Generate distinct colors with good visibility
+        # Using HSV color space for better distribution
+        hue = i / num_colors
+        saturation = 0.8 + np.random.random() * 0.2
+        value = 0.8 + np.random.random() * 0.2
+
+        # Convert HSV to BGR (OpenCV uses BGR)
+        hsv_color = np.array(
+            [[[hue * 180, saturation * 255, value * 255]]], dtype=np.uint8
+        )
+        bgr_color = cv2.cvtColor(hsv_color, cv2.COLOR_HSV2BGR)[0][0]
+
+        # Store as (B, G, R) tuple
+        colors[i] = (int(bgr_color[0]), int(bgr_color[1]), int(bgr_color[2]))
+
+    return colors