feat: Add native YOLO dataset import and in-place editing

anylabeling/views/labeling/label_widget.py

@@ -25,6 +25,15 @@
 from anylabeling.styles import AppTheme
 from anylabeling.views.labeling import utils
 from anylabeling.views.labeling.label_file import LabelFile, LabelFileError
+from anylabeling.views.labeling.yolo_io import (
+    find_dataset_config,
+    read_classes_txt,
+    read_dataset_yaml,
+    read_yolo_label,
+    write_yolo_label,
+    resolve_yolo_label_path,
+    update_dataset_config,
+)
 from anylabeling.views.labeling.logger import logger
 from anylabeling.views.labeling.shape import Shape
 from anylabeling.views.labeling.widgets import (
@@ -1196,6 +1205,14 @@ def __init__(
         self.output_file = output_file
         self.output_dir = output_dir
 
+        # YOLO-format annotation state
+        self._yolo_label_path = None
+        self._yolo_label_to_id = {}
+        self._yolo_id_to_label = {}
+        self._yolo_config_path = None
+        self._yolo_config_type = None
+        self._yolo_original_label_to_id = {}
+
         # Application state.
         self.image = QtGui.QImage()
         self.image_path = None
@@ -1399,6 +1416,7 @@ def reset_state(self):
         self.image_data = None
         self.label_file = None
         self.other_data = {}
+        self._yolo_label_path = None
         self.canvas.reset_state()
 
     def current_item(self):
@@ -1837,6 +1855,50 @@ def format_shape(s):
             key = item.text()
             flag = item.checkState() == Qt.CheckState.Checked
             flags[key] = flag
+
+        # Save to YOLO .txt when in YOLO mode
+        if self._yolo_label_path is not None:
+            try:
+                write_yolo_label(
+                    self._yolo_label_path,
+                    shapes,
+                    self.image.width(),
+                    self.image.height(),
+                    self._yolo_label_to_id,
+                )
+                self.label_file = None
+                if (
+                    self._yolo_config_path
+                    and self._yolo_label_to_id.keys()
+                    != self._yolo_original_label_to_id.keys()
+                ):
+                    update_dataset_config(
+                        self._yolo_config_path,
+                        self._yolo_config_type,
+                        self._yolo_label_to_id,
+                    )
+                    self._yolo_original_label_to_id = dict(
+                        self._yolo_label_to_id
+                    )
+                    self._yolo_id_to_label = {
+                        v: k for k, v in self._yolo_label_to_id.items()
+                    }
+            except Exception:
+                logger.warning(
+                    "Failed to save YOLO labels to %s",
+                    self._yolo_label_path,
+                    exc_info=True,
+                )
+                return False
+            items = self.file_list_widget.findItems(
+                self.image_path, Qt.MatchFlag.MatchExactly
+            )
+            if len(items) > 0:
+                if len(items) != 1:
+                    raise RuntimeError("There are duplicate files.")
+                items[0].setCheckState(Qt.CheckState.Checked)
+            return True
+
         try:
             image_path = osp.relpath(self.image_path, osp.dirname(filename))
             image_data = self.image_data if self._config["store_data"] else None
@@ -2163,6 +2225,21 @@ def load_file(self, filename=None):  # noqa: C901
             if self.image_data:
                 self.image_path = filename
             self.label_file = None
+        self._yolo_label_path = (
+            resolve_yolo_label_path(filename)
+            if self.label_file is None else None
+        )
+        # For a new image in a YOLO dataset (no .txt yet), construct
+        # the expected .txt path so annotations save in YOLO format.
+        if self._yolo_label_path is None and self._yolo_config_path is not None:
+            stem = osp.splitext(osp.basename(filename))[0]
+            img_dir = osp.dirname(filename)
+            candidate = osp.join(img_dir, stem + ".txt")
+            parent = osp.dirname(img_dir)
+            labels_dir = osp.join(parent, "labels")
+            if osp.isdir(labels_dir):
+                candidate = osp.join(labels_dir, stem + ".txt")
+            self._yolo_label_path = candidate
         image = QtGui.QImage.fromData(self.image_data)
 
         if image.isNull():
@@ -2185,6 +2262,30 @@ def load_file(self, filename=None):  # noqa: C901
             prev_shapes = self.canvas.shapes
         self.canvas.load_pixmap(QtGui.QPixmap.fromImage(image))
         flags = dict.fromkeys(self._config["flags"] or [], False)
+
+        # Try loading YOLO .txt annotations when no .json was found
+        if self._yolo_label_path is not None:
+            try:
+                if self._yolo_id_to_label and not self._yolo_label_to_id:
+                    self._yolo_label_to_id = {
+                        v: k for k, v in self._yolo_id_to_label.items()
+                    }
+                yolo_shapes = read_yolo_label(
+                    self._yolo_label_path,
+                    self.image.width(),
+                    self.image.height(),
+                    self._yolo_id_to_label,
+                    self._yolo_label_to_id,
+                )
+                if yolo_shapes:
+                    self.load_labels(yolo_shapes)
+            except Exception:
+                logger.warning(
+                    "Failed to load YOLO labels from %s",
+                    self._yolo_label_path,
+                    exc_info=True,
+                )
+
         if self.label_file:
             self.load_labels(self.label_file.shapes)
             if self.label_file.flags is not None:
@@ -2521,6 +2622,8 @@ def close_file(self, _value=False):
         self.actions.save_as.setEnabled(False)
 
     def get_label_file(self):
+        if self._yolo_label_path is not None:
+            return self._yolo_label_path
         if self.filename.lower().endswith(".json"):
             label_file = self.filename
         else:
@@ -2693,6 +2796,8 @@ def import_dropped_image_files(self, image_files):
             item.setFlags(Qt.ItemFlag.ItemIsEnabled | Qt.ItemFlag.ItemIsSelectable)
             if QtCore.QFile.exists(label_file) and LabelFile.is_label_file(label_file):
                 item.setCheckState(Qt.CheckState.Checked)
+            elif resolve_yolo_label_path(file) is not None:
+                item.setCheckState(Qt.CheckState.Checked)
             else:
                 item.setCheckState(Qt.CheckState.Unchecked)
             self.file_list_widget.addItem(item)
@@ -2710,6 +2815,40 @@ def import_image_folder(self, dirpath, pattern=None, load=True):
         if not self.may_continue() or not dirpath:
             return
 
+        # Detect YOLO dataset — look for data.yaml or classes.txt nearby
+        self._yolo_label_path = None
+        self._yolo_label_to_id = {}
+        self._yolo_id_to_label = {}
+        self._yolo_config_path = None
+        self._yolo_config_type = None
+        self._yolo_original_label_to_id = {}
+        config_path, config_type = find_dataset_config(dirpath)
+        if config_path:
+            if config_type == "yaml":
+                _, _, id_to_label = read_dataset_yaml(config_path)
+            else:
+                _, _, id_to_label = read_classes_txt(config_path)
+            if id_to_label:
+                self._yolo_id_to_label = id_to_label
+                self._yolo_label_to_id = {
+                    v: k for k, v in id_to_label.items()
+                }
+                self._yolo_original_label_to_id = dict(
+                    self._yolo_label_to_id
+                )
+                self._yolo_config_path = config_path
+                self._yolo_config_type = config_type
+                # Populate the label dialog and the unique label list
+                # with YOLO class names so users can select them from
+                # the sidebar and have them remembered for the next shape.
+                for name in self._yolo_label_to_id:
+                    self.label_dialog.add_label_history(name)
+                    if not self.unique_label_list.find_items_by_label(name):
+                        item = self.unique_label_list.create_item_from_label(name)
+                        self.unique_label_list.addItem(item)
+                        rgb = self._get_rgb_by_label(name)
+                        self.unique_label_list.set_item_label(item, name, rgb)
+
         self.last_open_dir = dirpath
         self.filename = None
         self.file_list_widget.clear()
@@ -2724,6 +2863,8 @@ def import_image_folder(self, dirpath, pattern=None, load=True):
             item.setFlags(Qt.ItemFlag.ItemIsEnabled | Qt.ItemFlag.ItemIsSelectable)
             if QtCore.QFile.exists(label_file) and LabelFile.is_label_file(label_file):
                 item.setCheckState(Qt.CheckState.Checked)
+            elif resolve_yolo_label_path(filename) is not None:
+                item.setCheckState(Qt.CheckState.Checked)
             else:
                 item.setCheckState(Qt.CheckState.Unchecked)
             self.file_list_widget.addItem(item)