imagebaker 0.0.50__py3-none-any.whl → 0.0.53__py3-none-any.whl

imagebaker/core/configs/configs.py

@@ -74,6 +74,22 @@ class BaseConfig(BaseModel):
             logger.info(f"Created assets folder at {asset_dir}")
         return asset_dir
 
+    @property
+    def cache_dir(self):
+        cache_dir = self.project_dir / ".imagebaker" / "cache"
+        if not cache_dir.exists():
+            cache_dir.mkdir(parents=True)
+            logger.info(f"Created cache folder at {cache_dir}")
+        return cache_dir
+
+    @property
+    def bake_dir(self):
+        bake_dir = self.project_dir / ".imagebaker" / "bake"
+        if not bake_dir.exists():
+            bake_dir.mkdir(parents=True)
+            logger.info(f"Created bake folder at {bake_dir}")
+        return bake_dir
+
     class Config:
         arbitrary_types_allowed = True
 
@@ -96,6 +112,7 @@ class LayerConfig(BaseConfig):
     )
     # whether to search image in subfolders as well
     full_search: bool = False
+    cleanup_on_exit: bool = False
 
     def get_label_color(self, label):
         for lbl in self.predefined_labels:

imagebaker/core/defs/defs.py

@@ -78,6 +78,7 @@ class LayerState:
     drawing_states: list[DrawingState] = field(default_factory=list)
     edge_opacity: int = 100
     edge_width: int = 10
+    caption: str = ""
 
     def copy(self):
         return LayerState(
@@ -109,6 +110,7 @@ class LayerState:
             ],
             edge_opacity=self.edge_opacity,
             edge_width=self.edge_width,
+            caption=self.caption,
         )
 
 
@@ -138,6 +140,7 @@ class Annotation:
     file_path: Path = field(default_factory=lambda: Path("Runtime"))
     is_model_generated: bool = False
     model_name: str = None
+    caption: str = ""
 
     def copy(self):
         ann = Annotation(
@@ -156,6 +159,7 @@ class Annotation:
             file_path=self.file_path,
             is_model_generated=self.is_model_generated,
             model_name=self.model_name,
+            caption=self.caption,
         )
         ann.is_selected = False
         return ann
@@ -200,6 +204,7 @@ class Annotation:
                 "file_path": str(annotation.file_path),
                 "is_model_generated": annotation.is_model_generated,
                 "model_name": annotation.model_name,
+                "caption": annotation.caption,
             }
             annotations_dict.append(data)
 
@@ -210,6 +215,10 @@ class Annotation:
     def load_from_json(path: str):
         import json
 
+        # if path does not exist, return empty list
+        if not Path(path).exists():
+            return []
+
         with open(path, "r") as f:
             data = json.load(f)
 
@@ -230,6 +239,7 @@ class Annotation:
                 file_path=Path(d.get("file_path", "Runtime")),
                 is_model_generated=d.get("is_model_generated", False),
                 model_name=d.get("model_name", None),
+                caption=d.get("caption", ""),
             )
 
             # Handle points safely

imagebaker/layers/annotable_layer.py

@@ -24,6 +24,7 @@ from PySide6.QtGui import (
     QWheelEvent,
     QMouseEvent,
     QKeyEvent,
+    QFont,
 )
 from PySide6.QtWidgets import (
     QApplication,
@@ -42,6 +43,7 @@ class AnnotableLayer(BaseLayer):
     annotationCleared = Signal()
     annotationMoved = Signal()
     layersChanged = Signal()
+    labelUpdated = Signal(tuple)
 
     def __init__(self, parent, config: LayerConfig, canvas_config: CanvasConfig):
         super().__init__(parent, config)
@@ -67,14 +69,23 @@ class AnnotableLayer(BaseLayer):
         self.annotationCleared.emit()
         self.update()
 
+    def toggle_annotation_visibility(self):
+        """Toggle visibility of all annotations."""
+        selected_annotation = self._get_selected_annotation()
+        if selected_annotation is not None:
+            selected_annotation.visible = not selected_annotation.visible
+            self.annotationUpdated.emit(selected_annotation)
+            self.update()
+
     def handle_key_press(self, event: QKeyEvent):
         # Handle Ctrl key for panning
         if event.key() == Qt.Key_Control:
-            if (
-                self.mouse_mode != MouseMode.POLYGON
-            ):  # Only activate pan mode when not drawing polygons
+            # disable this for now to allow pannings
+            # if (
+            #     self.mouse_mode != MouseMode.POLYGON
+            # ):  # Only activate pan mode when not drawing polygons
 
-                self.mouse_mode = MouseMode.PAN
+            self.mouse_mode = MouseMode.PAN
 
         # Handle Ctrl+C for copy
         if event.modifiers() & Qt.ControlModifier and event.key() == Qt.Key_C:
@@ -257,7 +268,7 @@ class AnnotableLayer(BaseLayer):
 
         # Draw labels
         if annotation.is_complete and annotation.label:
-            painter.save()
+            # painter.save()
             label_pos = self.get_label_position(annotation)
             text = annotation.label
 
@@ -297,8 +308,58 @@ class AnnotableLayer(BaseLayer):
             # Draw text
             painter.setPen(Qt.white)
             painter.drawText(bg_rect, Qt.AlignCenter, text)
-            painter.restore()
-            self.label_rects.append((bg_rect, annotation))
+
+            # now if the annotations has caption,
+            # draw the caption below the label in itallic
+            # just below the label, in a font lighter than a label
+
+            if annotation.caption:
+                caption_font = painter.font()
+                # Draw a background rectangle for the caption to ensure visibility
+                caption_text = annotation.caption
+                caption_font.setItalic(True)
+                caption_font.setPixelSize(
+                    self.config.selected_draw_config.label_font_size * self.scale + 2
+                )
+                caption_font.setWeight(QFont.Light)
+                painter.setFont(caption_font)
+                metrics = painter.fontMetrics()
+                text_width = metrics.horizontalAdvance(caption_text)
+                text_height = metrics.height()
+                # Draw background rectangle for caption
+                caption_rect = QRectF(
+                    widget_pos.x() - text_width / 2 - 2,
+                    widget_pos.y() + text_height / 2 - 2,
+                    text_width + 4,
+                    text_height + 4,
+                )
+                painter.setPen(Qt.NoPen)
+                painter.setBrush(
+                    self.config.normal_draw_config.label_font_background_color
+                )
+                painter.drawRect(caption_rect)
+                # Draw caption text
+                painter.setPen(Qt.white)
+                painter.drawText(caption_rect, Qt.AlignCenter, caption_text)
+                caption_font.setItalic(True)
+                caption_font.setPixelSize(
+                    self.config.selected_draw_config.label_font_size * self.scale + 2
+                )
+                caption_font.setWeight(QFont.Light)
+                painter.setFont(caption_font)
+                metrics = painter.fontMetrics()
+                text_width = metrics.horizontalAdvance(annotation.caption)
+                text_height = metrics.height()
+                painter.setPen(Qt.white)
+                painter.setBrush(Qt.gray)
+
+                caption_rect = QRectF(
+                    widget_pos.x() - text_width / 2 - 2,
+                    widget_pos.y() + text_height / 2 - 2,
+                    text_width + 4,
+                    text_height + 4,
+                )
+                painter.drawText(caption_rect, Qt.AlignCenter, annotation.caption)
 
         painter.restore()
 
@@ -397,7 +458,14 @@ class AnnotableLayer(BaseLayer):
                 MouseMode.ZOOM_IN,
                 MouseMode.ZOOM_OUT,
             ]:
-                self.mouse_mode = MouseMode.IDLE
+                if self.current_annotation:
+                    if (
+                        self.current_annotation.polygon
+                        and not self.current_annotation.is_complete
+                    ):
+                        self.mouse_mode = MouseMode.POLYGON
+                else:
+                    self.mouse_mode = MouseMode.IDLE
 
         # Clean up transformation state
         if hasattr(self, "selected_annotation"):
@@ -534,13 +602,11 @@ class AnnotableLayer(BaseLayer):
                     self.current_annotation.polygon = QPolygonF(
                         [p for p in self.current_annotation.polygon][:-1]
                     )
-                    self.update()
 
                 # If the polygon is now empty, reset to idle mode
                 if len(self.current_annotation.polygon) == 0:
                     self.current_annotation = None
                     self.mouse_mode = MouseMode.IDLE
-                    self.update()
 
             # If not drawing a polygon, go to idle mode
             if not self.current_annotation:
@@ -548,7 +614,7 @@ class AnnotableLayer(BaseLayer):
                 for ann in self.annotations:
                     ann.selected = False
                     self.annotationUpdated.emit(ann)
-                self.update()
+            self.update()
 
         # If left-clicked
         if event.button() == Qt.LeftButton:
@@ -603,6 +669,17 @@ class AnnotableLayer(BaseLayer):
             elif self.mouse_mode == MouseMode.POLYGON:
                 # If not double-click
                 if not self.current_annotation:
+                    # if this point is equal to the last point of the previous polygon, then ignore it
+                    if len(self.annotations) > 0:
+                        last_polygon = self.annotations[-1].polygon
+                        if last_polygon:
+                            last_point = last_polygon[-1]
+                            if last_point == clamped_pos:
+                                logger.info(
+                                    "Ignoring point, same as last polygon point"
+                                )
+                                return
+
                     self.current_annotation = Annotation(
                         file_path=self.file_path,
                         annotation_id=len(self.annotations),
@@ -751,8 +828,8 @@ class AnnotableLayer(BaseLayer):
                 and len(self.current_annotation.polygon) >= 3
             ):
                 self.current_annotation.is_complete = True
+
                 self.finalize_annotation()
-                self.annotationAdded.emit(self.current_annotation)
                 self.current_annotation = None
 
                 return
@@ -794,6 +871,7 @@ class AnnotableLayer(BaseLayer):
             self, "Edit Label", "Enter new label:", text=annotation.label
         )
         if ok and new_label:
+            self.labelUpdated.emit((annotation.label, new_label))
             annotation.label = new_label
             self.annotationUpdated.emit(annotation)  # Emit signal
             self.update()
@@ -811,8 +889,6 @@ class AnnotableLayer(BaseLayer):
                 self.current_annotation
             )
             self.annotationAdded.emit(self.current_annotation)
-            self.current_annotation = None
-            self.update()
         else:
             # Show custom label dialog
             label, ok = QInputDialog.getText(self, "Label", "Enter label name:")
@@ -827,8 +903,8 @@ class AnnotableLayer(BaseLayer):
                     )
                     self.annotationAdded.emit(self.current_annotation)
                     self.current_annotation.annotation_id = len(self.annotations)
-                    self.current_annotation = None
-                    self.update()
+        self.current_annotation = None
+        self.update()
 
     # in update, update cursor
 
@@ -915,6 +991,7 @@ class AnnotableLayer(BaseLayer):
         new_layer.layer_name = (
             f"{annotation.label} {annotation.annotation_id} {annotation.annotator}"
         )
+        new_layer.caption = annotation.caption
 
         new_layer._apply_edge_opacity()
         new_layer.update()

imagebaker/layers/base_layer.py

@@ -240,7 +240,7 @@ class BaseLayer(QWidget):
             for step, state in enumerate(intermediate_states):
                 step += self.current_step
 
-                logger.info(f"Saving state {step} for layer {layer.layer_id}")
+                logger.info(f"Saving state {step} for layer {layer.layer_id}: {state}")
                 state.selected = False
                 if step not in curr_states:
                     curr_states[step] = []
@@ -853,3 +853,11 @@ class BaseLayer(QWidget):
     @edge_width.setter
     def edge_width(self, value: int):
         self.layer_state.edge_width = value
+
+    @property
+    def caption(self) -> str:
+        return self.layer_state.caption
+
+    @caption.setter
+    def caption(self, value: str):
+        self.layer_state.caption = value

imagebaker/layers/canvas_layer.py

@@ -27,7 +27,6 @@ from PySide6.QtGui import (
     QMouseEvent,
     QKeyEvent,
     QTransform,
-    QImage,
 )
 from PySide6.QtWidgets import (
     QApplication,
@@ -127,7 +126,6 @@ class CanvasLayer(BaseLayer):
 
     ## Helper functions ##
     def handle_key_press(self, event: QKeyEvent):
-
         # Handle Delete key
         if event.key() == Qt.Key_Delete:
             self._delete_layer()
@@ -149,6 +147,15 @@ class CanvasLayer(BaseLayer):
         if event.modifiers() & Qt.ControlModifier and event.key() == Qt.Key_V:
             self._paste_layer()
             return  # Important: return after handling
+        # if event.key() == Qt.Key_H:
+        #     if self.selected_layer:
+        #         # Toggle visibility of the selected layer
+        #         logger.info(
+        #             f"Toggling visibility of layer: {self.selected_layer.layer_name}"
+        #         )
+        #         self.selected_layer.visible = not self.selected_layer.visible
+        #         self.selected_layer.update()
+        #     return  # Important: return after handling
 
     def paint_layer(self, painter: QPainter):
         """
@@ -362,7 +369,16 @@ class CanvasLayer(BaseLayer):
         """
         """Add a new drawing state."""
         self.selected_layer = self._get_selected_layer()
-        layer = self.selected_layer if self.selected_layer else self
+        layer = self.selected_layer
+        if not layer:
+            logger.debug("No layer selected for drawing.")
+            # show popup message window that closes in 2 seconds
+            QMessageBox.information(
+                self.parent(),
+                "No Layer Selected",
+                "Please select a layer to draw on.",
+            )
+            return
 
         # Convert the position to be relative to the layer
         relative_pos = pos - layer.position
@@ -384,8 +400,14 @@ class CanvasLayer(BaseLayer):
             )
             layer.layer_state.drawing_states.append(drawing_state)
             self._last_draw_point = relative_pos  # Update the last draw point
-            # logger.debug(f"Added drawing state at position: {relative_pos}")
-
+            logger.debug(
+                f"Added drawing state at position: {relative_pos} to layer {layer.layer_name}"
+            )
+        else:
+            logger.debug(
+                f"Mouse mode {self.mouse_mode} does not support drawing states."
+            )
+            return
         self.update()  # Refresh the canvas to show the new drawing
 
     def handle_wheel(self, event: QWheelEvent):
@@ -822,6 +844,7 @@ class CanvasLayer(BaseLayer):
         filename = self.config.export_folder / f"{filename}.png"
         logger.info(f"Exporting baked image to {filename}")
         self.states = {0: [layer.layer_state for layer in self.layers]}
+        logger.debug(f"Exporting states: {self.states}")
 
         self.loading_dialog = QProgressDialog(
             "Baking Please wait...", "Cancel", 0, 0, self.parentWidget()
@@ -1030,7 +1053,7 @@ class CanvasLayer(BaseLayer):
                     logger.info(f"Saved baked image to {filename}")
                     if self.config.write_annotations:
                         image = qpixmap_to_numpy(image.copy())
-                        image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGR)
+                        image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA)
                         drawn = draw_annotations(image, annotations)
                         write_to = filename.parent / f"annotated_{filename.name}"
 

imagebaker/list_views/annotation_list.py

@@ -52,6 +52,7 @@ class AnnotationList(QDockWidget):
         if self.layer is None:
             return
         for idx, ann in enumerate(self.layer.annotations):
+
             item = QListWidgetItem(self.list_widget)
 
             # Create container widget
@@ -200,5 +201,6 @@ class AnnotationList(QDockWidget):
             self.layer.paste_annotation()
         elif event.key() == Qt.Key_Delete:
             self.delete_annotation(self.layer.selected_annotation_index)
-        else:
-            self.parentWidget().keyPressEvent(event)
+        elif event.key() == Qt.Key_H:
+            self.layer.toggle_annotation_visibility()
+        # self.parentWidget().keyPressEvent(event)

imagebaker/list_views/image_list.py

@@ -19,6 +19,7 @@ from pathlib import Path
 
 class ImageListPanel(QDockWidget):
     imageSelected = Signal(Path)
+    activeImageEntries = Signal(list)
 
     def __init__(
         self,
@@ -26,6 +27,7 @@ class ImageListPanel(QDockWidget):
         processed_images: set[Path],
         parent=None,
         max_name_length=15,
+        images_per_page=10,
     ):
         """
         :param image_entries: List of image paths to display.
@@ -35,7 +37,7 @@ class ImageListPanel(QDockWidget):
         self.image_entries: list["ImageEntry"] = image_entries
         self.processed_images = processed_images
         self.current_page = 0
-        self.images_per_page = 10
+        self.images_per_page = images_per_page
         self.max_name_length = max_name_length
         self.init_ui()
 
@@ -90,7 +92,20 @@ class ImageListPanel(QDockWidget):
         """Update the image list with image paths and baked results."""
         self.list_widget.clear()
 
-        for idx, image_entry in enumerate(image_entries):
+        # Calculate the range of images to display for the current page
+        start_index = self.current_page * self.images_per_page
+        end_index = min(start_index + self.images_per_page, len(image_entries))
+
+        # Update the pagination label
+        self.pagination_label.setText(
+            f"Showing {start_index + 1} to {end_index} of {len(image_entries)}"
+        )
+
+        # Display only the images for the current page
+        active_image_entries = []
+        for idx, image_entry in enumerate(
+            image_entries[start_index:end_index], start=start_index + 1
+        ):
             item_widget = QWidget()
             item_layout = QHBoxLayout(item_widget)
             item_layout.setContentsMargins(5, 5, 5, 5)
@@ -101,7 +116,7 @@ class ImageListPanel(QDockWidget):
                 thumbnail_pixmap = (
                     image_entry.data.get_thumbnail()
                 )  # Baked result thumbnail
-                name_label_text = f"Baked Result {idx + 1}"
+                name_label_text = f"Baked Result {idx}"
             else:
                 thumbnail_pixmap = QPixmap(str(image_entry.data)).scaled(
                     50, 50, Qt.KeepAspectRatio, Qt.SmoothTransformation
@@ -126,11 +141,14 @@ class ImageListPanel(QDockWidget):
 
             # Store metadata for the image
             list_item.setData(Qt.UserRole, image_entry)
-        self.pagination_label.setText(
-            f"Showing {self.current_page * self.images_per_page + 1} to "
-            f"{min((self.current_page + 1) * self.images_per_page, len(image_entries))} "
-            f"of {len(image_entries)}"
-        )
+            active_image_entries.append(image_entry)
+
+        self.activeImageEntries.emit(active_image_entries)
+        # also set first item as selected if there are any items
+        if self.list_widget.count() > 0:
+            self.list_widget.setCurrentRow(0)
+            self.list_widget.setFocus()
+
         self.update()
 
     def handle_item_clicked(self, item: QListWidgetItem):

imagebaker/list_views/layer_settings.py

@@ -8,6 +8,7 @@ from PySide6.QtWidgets import (
     QSizePolicy,
     QDockWidget,
     QSlider,
+    QLineEdit,
 )
 
 from imagebaker.core.defs import LayerState
@@ -82,6 +83,32 @@ class LayerSettings(QDockWidget):
         )
         self.main_layout.addWidget(self.edge_width_slider["widget"])
 
+        # Caption input
+        caption_layout = QHBoxLayout()
+        caption_label = QLabel("Caption:")
+        caption_label.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed)
+        self.caption_input = QLineEdit()
+        self.caption_input.setPlaceholderText("Enter caption...")
+        self.caption_input.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
+        caption_layout.addWidget(caption_label)
+        caption_layout.addWidget(self.caption_input)
+        self.main_layout.addLayout(caption_layout)
+
+        # Set initial value if layer is selected
+        if self.selected_layer:
+            self.caption_input.setText(self.selected_layer.caption)
+
+        # Update layer caption on editing finished
+        def update_caption():
+            if self.selected_layer:
+                logger.info(
+                    f"Updating caption for layer {self.selected_layer.layer_name} to {self.caption_input.text()}"
+                )
+                self.selected_layer.caption = self.caption_input.text()
+                self.selected_layer.update()
+
+        self.caption_input.editingFinished.connect(update_caption)
+
         # Add stretch to push content to the top
         self.main_layout.addStretch()
 
@@ -177,6 +204,7 @@ class LayerSettings(QDockWidget):
             edge_opacity=self.selected_layer.edge_opacity,
             edge_width=self.selected_layer.edge_width,
             visible=self.selected_layer.visible,
+            caption=self.selected_layer.caption,
         )
         logger.info(f"Storing state {bake_settings}")
         self.messageSignal.emit(f"Stored state for {bake_settings.layer_name}")
@@ -240,6 +268,7 @@ class LayerSettings(QDockWidget):
                 self.edge_width_slider["slider"].setValue(
                     int(self.selected_layer.edge_width)
                 )
+                self.caption_input.setText(self.selected_layer.caption)
             else:
                 self.widget.setEnabled(False)
                 self.layer_name_label.setText("No BaseLayer")

imagebaker/tabs/baker_tab.py

@@ -18,7 +18,6 @@ from PySide6.QtWidgets import (
     QSlider,
     QLabel,
     QSpinBox,
-    QComboBox,
 )
 from collections import deque
 
@@ -485,5 +484,54 @@ class BakerTab(QWidget):
             self.layer_list.update_list()
             self.messageSignal.emit(f"Added new layer: {new_layer.layer_name}")
 
+        if event.key() == Qt.Key_H:
+            # if pressed h key, then toggle visibility of layer
+            selected_layer = self.current_canvas._get_selected_layer()
+            if selected_layer:
+                selected_layer.visible = not selected_layer.visible
+                selected_layer.update()
+                self.current_canvas.update()
+                self.layer_settings.update_sliders()
+                self.messageSignal.emit(
+                    f"Toggled visibility of annotations in layer: {selected_layer.layer_name}"
+                )
+        # if pressed W key, move the order of the selected layer up
+        # make it circular, if the first layer is selected, then move it to the last
+        if event.key() == Qt.Key_W:
+            selected_layer = self.current_canvas._get_selected_layer()
+            if selected_layer:
+                index = self.current_canvas.layers.index(selected_layer)
+                if index > 0:
+                    # Move the layer up
+                    self.current_canvas.layers.insert(index - 1, selected_layer)
+                    del self.current_canvas.layers[index + 1]
+                else:
+                    # Move the first layer to the end
+                    self.current_canvas.layers.append(selected_layer)
+                    del self.current_canvas.layers[0]
+                self.current_canvas.update()
+                self.layer_list.update_list()
+                self.messageSignal.emit(
+                    f"Moved layer {selected_layer.layer_name} up in order."
+                )
+        # if pressed S key, move the order of the selected layer down
+        # make it circular, if the last layer is selected, then move it to the first
+        if event.key() == Qt.Key_S:
+            selected_layer = self.current_canvas._get_selected_layer()
+            if selected_layer:
+                index = self.current_canvas.layers.index(selected_layer)
+                if index < len(self.current_canvas.layers) - 1:
+                    # Move the layer down
+                    self.current_canvas.layers.insert(index + 2, selected_layer)
+                    del self.current_canvas.layers[index]
+                else:
+                    # Move the last layer to the beginning
+                    self.current_canvas.layers.insert(0, selected_layer)
+                    del self.current_canvas.layers[-1]
+                self.current_canvas.update()
+                self.layer_list.update_list()
+                self.messageSignal.emit(
+                    f"Moved layer {selected_layer.layer_name} down in order."
+                )
         self.update()
         return super().keyPressEvent(event)