From e3ad5127fb062a8dc2d751587baee97801f19cfe Mon Sep 17 00:00:00 2001 From: Leon Kraim Date: Thu, 9 Jul 2026 15:24:34 +0200 Subject: [PATCH 1/6] Add native YOLO dataset import and in-place editing Allows opening a YOLO-format dataset directory via Open Dir (Ctrl+U) and editing annotations in-place without creating AnyLabeling .json files. - New anylabeling/views/labeling/yolo_io.py with utilities: - find_dataset_yaml() walks up 5 levels looking for data.yaml - read_dataset_yaml() parses nc + names into id-to-label mapping - read_yolo_label() parses .txt files into AnyLabeling shapes - write_yolo_label() writes shapes back to YOLO .txt format - resolve_yolo_label_path() finds sibling labels/ directory - Modified label_widget.py: - import_image_folder() detects data.yaml, populates class map - load_file() falls back to .txt when no .json exists - save_labels() writes .txt when source was YOLO format - reset_state() clears YOLO path - get_label_file() returns .txt path in YOLO mode - import_dropped_image_files() checks for .txt checkmarks Priority: .json > .txt (same dir) > .txt (sibling labels/) No regression on existing JSON-only workflows. --- anylabeling/views/labeling/label_widget.py | 73 ++++++++ anylabeling/views/labeling/yolo_io.py | 188 +++++++++++++++++++++ 2 files changed, 261 insertions(+) create mode 100644 anylabeling/views/labeling/yolo_io.py diff --git a/anylabeling/views/labeling/label_widget.py b/anylabeling/views/labeling/label_widget.py index 6816dcf..3bfd4fd 100644 --- a/anylabeling/views/labeling/label_widget.py +++ b/anylabeling/views/labeling/label_widget.py @@ -25,6 +25,13 @@ 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_yaml, + read_dataset_yaml, + read_yolo_label, + write_yolo_label, + resolve_yolo_label_path, +) from anylabeling.views.labeling.logger import logger from anylabeling.views.labeling.shape import Shape from anylabeling.views.labeling.widgets import ( @@ -1196,6 +1203,11 @@ 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 = {} + # Application state. self.image = QtGui.QImage() self.image_path = None @@ -1399,6 +1411,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 +1850,26 @@ 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: + write_yolo_label( + self._yolo_label_path, + shapes, + self.image.width(), + self.image.height(), + self._yolo_label_to_id, + ) + self.label_file = None + 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 @@ -2185,6 +2218,27 @@ 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.label_file is None and self._yolo_label_path is None: + yolo_txt_path = resolve_yolo_label_path(filename) + if yolo_txt_path is not None: + self._yolo_label_path = yolo_txt_path + # Ensure label_to_id is in sync with id_to_label + 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( + yolo_txt_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) + if self.label_file: self.load_labels(self.label_file.shapes) if self.label_file.flags is not None: @@ -2521,6 +2575,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 +2749,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 +2768,19 @@ 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 nearby + self._yolo_label_path = None + self._yolo_label_to_id = {} + self._yolo_id_to_label = {} + yaml_path = find_dataset_yaml(dirpath) + if yaml_path: + _, _, id_to_label = read_dataset_yaml(yaml_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.last_open_dir = dirpath self.filename = None self.file_list_widget.clear() @@ -2724,6 +2795,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) diff --git a/anylabeling/views/labeling/yolo_io.py b/anylabeling/views/labeling/yolo_io.py new file mode 100644 index 0000000..ed85ca4 --- /dev/null +++ b/anylabeling/views/labeling/yolo_io.py @@ -0,0 +1,188 @@ +"""Utilities for reading and writing YOLO-format annotations.""" + +import os +import os.path as osp + + +def find_dataset_yaml(dirpath): + """Walk up from dirpath looking for a data.yaml file. + + Searches up to 5 parent directories. Returns the full path to + data.yaml if found, None otherwise. + """ + current = osp.abspath(dirpath) + for _ in range(5): + candidate = osp.join(current, "data.yaml") + if osp.exists(candidate): + return candidate + parent = osp.dirname(current) + if parent == current: + break + current = parent + return None + + +def read_dataset_yaml(yaml_path): + """Parse a YOLO data.yaml file and return class name information. + + Returns (nc, names_list, id_to_label_dict). Any of the three + return values may be empty / 0 when the key is absent from the + YAML file. + """ + import yaml + + with open(yaml_path, "r", encoding="utf-8") as f: + data = yaml.safe_load(f) + + nc = data.get("nc", 0) + raw_names = data.get("names", []) + + if isinstance(raw_names, list): + id_to_label = {i: name for i, name in enumerate(raw_names)} + elif isinstance(raw_names, dict): + id_to_label = {int(k): v for k, v in raw_names.items()} + else: + id_to_label = {} + return nc, raw_names, id_to_label + + +def read_yolo_label(txt_path, img_w, img_h, id_to_label, label_to_id=None): + """Parse a YOLO-format .txt label file into AnyLabeling shape dicts. + + Each line in the file is expected to follow the YOLO detection + format:: + + + + Coordinates are normalised (0-1). Returns a list of dicts ready + to be passed to ``LabelingWidget.load_labels()``. + + When *label_to_id* is provided the dict will be updated in-place + so that every label string that appears in the file maps back to + its original class ID, ensuring round-trip fidelity. + """ + shapes = [] + if not osp.exists(txt_path): + return shapes + + with open(txt_path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line or line.startswith("#"): + continue + parts = line.split() + if len(parts) < 5: + continue + + class_id = int(parts[0]) + x_center = float(parts[1]) + y_center = float(parts[2]) + width = float(parts[3]) + height = float(parts[4]) + + # Normalised → absolute pixel coords + w_abs = width * img_w + h_abs = height * img_h + cx_abs = x_center * img_w + cy_abs = y_center * img_h + + x1 = cx_abs - w_abs / 2.0 + y1 = cy_abs - h_abs / 2.0 + x2 = cx_abs + w_abs / 2.0 + y2 = cy_abs + h_abs / 2.0 + + label = id_to_label.get(class_id) + if label is None: + label = f"class_{class_id}" + + if label_to_id is not None and label not in label_to_id: + label_to_id[label] = class_id + + shape = { + "label": label, + "text": "", + "points": [[x1, y1], [x2, y2]], + "shape_type": "rectangle", + "group_id": None, + "flags": {}, + "other_data": {}, + } + shapes.append(shape) + + return shapes + + +def write_yolo_label(txt_path, shapes, img_w, img_h, label_to_id): + """Write AnyLabeling shape dicts back to a YOLO-format .txt file. + + *label_to_id* is mutated in-place: labels that are not already + present receive the next unused integer class ID so that saving + never silently drops annotations. + """ + used_ids = set(label_to_id.values()) + next_id = 0 + if used_ids: + next_id = max(used_ids) + 1 + + lines = [] + for shape in shapes: + if shape["shape_type"] not in ("rectangle",): + continue + points = shape["points"] + if len(points) < 2: + continue + + label = shape["label"] + + if label not in label_to_id: + while next_id in used_ids: + next_id += 1 + label_to_id[label] = next_id + used_ids.add(next_id) + next_id += 1 + + class_id = label_to_id[label] + + # Absolute pixel coords → normalised YOLO format + x1, y1 = points[0] + x2, y2 = points[1] + + x_center = ((x1 + x2) / 2.0) / img_w + y_center = ((y1 + y2) / 2.0) / img_h + width = abs(x2 - x1) / img_w + height = abs(y2 - y1) / img_h + + lines.append( + f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}" + ) + + with open(txt_path, "w", encoding="utf-8") as f: + f.write("\n".join(lines)) + if lines: + f.write("\n") + + +def resolve_yolo_label_path(image_path): + """Locate the YOLO .txt label file that corresponds to *image_path*. + + Checks, in order: + 1. Same directory as the image (``.txt``) + 2. A sibling ``labels/`` directory (``../labels/.txt``) + + Returns the absolute path to the .txt file or ``None``. + """ + img_dir = osp.dirname(image_path) + base_stem = osp.splitext(osp.basename(image_path))[0] + + # 1. Side-by-side: same directory + candidate = osp.join(img_dir, base_stem + ".txt") + if osp.exists(candidate): + return candidate + + # 2. Standard YOLO layout: ../labels/.txt + parent = osp.dirname(img_dir) + candidate = osp.join(parent, "labels", base_stem + ".txt") + if osp.exists(candidate): + return candidate + + return None From a5fc5fa2562a975aad5183caf1759b8477d12341 Mon Sep 17 00:00:00 2001 From: Leon Kraim Date: Thu, 9 Jul 2026 17:43:57 +0200 Subject: [PATCH 2/6] Fix YOLO navigation bug, sync data.yaml on class changes, support classes.txt Fixes a bug where _yolo_label_path was set once and never reset, causing YOLO labels from the first image to persist when navigating to other images, and saving to overwrite the wrong .txt file. Now resolves the path per-image in load_file(). Adds data.yaml / classes.txt sync: - find_dataset_config() discovers both data.yaml and classes.txt - update_dataset_config() keeps config in sync when new classes are added in the UI (targeted regex for YAML, plain rewrite for txt) - Preserves all other YAML content (comments, train/val, roboflow) Adds error handling with try/except + logging on all YOLO I/O paths to prevent crashes from corrupt label files. --- anylabeling/views/labeling/label_widget.py | 84 +++++++++++++++----- anylabeling/views/labeling/yolo_io.py | 89 ++++++++++++++++++++++ 2 files changed, 154 insertions(+), 19 deletions(-) diff --git a/anylabeling/views/labeling/label_widget.py b/anylabeling/views/labeling/label_widget.py index 3bfd4fd..4601c25 100644 --- a/anylabeling/views/labeling/label_widget.py +++ b/anylabeling/views/labeling/label_widget.py @@ -26,11 +26,13 @@ from anylabeling.views.labeling import utils from anylabeling.views.labeling.label_file import LabelFile, LabelFileError from anylabeling.views.labeling.yolo_io import ( - find_dataset_yaml, + 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 @@ -1207,6 +1209,9 @@ def __init__( 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() @@ -1412,6 +1417,8 @@ def reset_state(self): self.label_file = None self.other_data = {} self._yolo_label_path = None + self._yolo_config_path = None + self._yolo_config_type = None self.canvas.reset_state() def current_item(self): @@ -1853,14 +1860,35 @@ def format_shape(s): # Save to YOLO .txt when in YOLO mode if self._yolo_label_path is not None: - write_yolo_label( - self._yolo_label_path, - shapes, - self.image.width(), - self.image.height(), - self._yolo_label_to_id, - ) - self.label_file = 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 + ) + 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 ) @@ -2196,6 +2224,10 @@ 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 + ) image = QtGui.QImage.fromData(self.image_data) if image.isNull(): @@ -2220,17 +2252,14 @@ def load_file(self, filename=None): # noqa: C901 flags = dict.fromkeys(self._config["flags"] or [], False) # Try loading YOLO .txt annotations when no .json was found - if self.label_file is None and self._yolo_label_path is None: - yolo_txt_path = resolve_yolo_label_path(filename) - if yolo_txt_path is not None: - self._yolo_label_path = yolo_txt_path - # Ensure label_to_id is in sync with id_to_label + 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( - yolo_txt_path, + self._yolo_label_path, self.image.width(), self.image.height(), self._yolo_id_to_label, @@ -2238,6 +2267,12 @@ def load_file(self, filename=None): # noqa: C901 ) 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) @@ -2768,18 +2803,29 @@ 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 nearby + # 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 = {} - yaml_path = find_dataset_yaml(dirpath) - if yaml_path: - _, _, id_to_label = read_dataset_yaml(yaml_path) + 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 self.last_open_dir = dirpath self.filename = None diff --git a/anylabeling/views/labeling/yolo_io.py b/anylabeling/views/labeling/yolo_io.py index ed85ca4..8802651 100644 --- a/anylabeling/views/labeling/yolo_io.py +++ b/anylabeling/views/labeling/yolo_io.py @@ -162,6 +162,95 @@ def write_yolo_label(txt_path, shapes, img_w, img_h, label_to_id): f.write("\n") +def find_dataset_config(dirpath): + """Walk up from dirpath looking for a YOLO dataset config file. + + Checks, in order: + 1. ``data.yaml`` (YOLO dataset descriptor, YAML format) + 2. ``classes.txt`` (one class name per line, plain text) + + Searches up to 5 parent directories. Returns ``(path, type)`` + where *type* is ``"yaml"`` or ``"txt"``, or ``(None, None)``. + """ + current = osp.abspath(dirpath) + for _ in range(5): + for filename, cfg_type in (("data.yaml", "yaml"), ("classes.txt", "txt")): + candidate = osp.join(current, filename) + if osp.exists(candidate): + return candidate, cfg_type + parent = osp.dirname(current) + if parent == current: + break + current = parent + return None, None + + +def read_classes_txt(txt_path): + """Parse a classes.txt file (one class name per line). + + Returns ``(nc, names_list, id_to_label_dict)``. + """ + names = [] + with open(txt_path, "r", encoding="utf-8") as f: + for line in f: + name = line.strip() + if name and not name.startswith("#"): + names.append(name) + nc = len(names) + id_to_label = {i: name for i, name in enumerate(names)} + return nc, names, id_to_label + + +def update_dataset_config(config_path, config_type, label_to_id): + """Write updated class names back to a YOLO dataset config file. + + * For ``"yaml"``: only the ``nc:`` and ``names:`` lines are + replaced in-place via regex — every other line (comments, + train/val paths, roboflow metadata, blank lines) is preserved + character-for-character. + + * For ``"txt"``: the file is rewritten with one class name per + line in class-ID order. + + *label_to_id* is a ``{name: class_id}`` mapping. + """ + if not label_to_id: + return + max_id = max(label_to_id.values()) + id_to_label = {} + for name, cid in label_to_id.items(): + id_to_label[cid] = name + ordered_names = [id_to_label.get(i, "") for i in range(max_id + 1)] + for i, name in enumerate(ordered_names): + if not name: + ordered_names[i] = f"class_{i}" + + if config_type == "yaml": + import re + + with open(config_path, "r", encoding="utf-8") as f: + content = f.read() + content = re.sub( + r"^nc:\s*\d+", + f"nc: {len(ordered_names)}", + content, + flags=re.MULTILINE, + ) + names_str = repr(ordered_names) + content = re.sub( + r"^names:\s*\[.*?\]|^names:\s*\n(\s+- .*\n?)*", + f"names: {names_str}", + content, + flags=re.MULTILINE | re.DOTALL, + ) + with open(config_path, "w", encoding="utf-8") as f: + f.write(content) + else: + with open(config_path, "w", encoding="utf-8") as f: + for name in ordered_names: + f.write(name + "\n") + + def resolve_yolo_label_path(image_path): """Locate the YOLO .txt label file that corresponds to *image_path*. From 3ae5a0cc2c9f0c454bbf0498839ffbfae88dd4e2 Mon Sep 17 00:00:00 2001 From: Leon Kraim Date: Thu, 9 Jul 2026 20:48:13 +0200 Subject: [PATCH 3/6] Support YOLO segmentation format (polygons) in import/export --- anylabeling/views/labeling/yolo_io.py | 126 ++++++--- tests/test_yolo_io.py | 386 ++++++++++++++++++++++++++ 2 files changed, 470 insertions(+), 42 deletions(-) create mode 100644 tests/test_yolo_io.py diff --git a/anylabeling/views/labeling/yolo_io.py b/anylabeling/views/labeling/yolo_io.py index 8802651..9d51fc0 100644 --- a/anylabeling/views/labeling/yolo_io.py +++ b/anylabeling/views/labeling/yolo_io.py @@ -49,11 +49,17 @@ def read_dataset_yaml(yaml_path): def read_yolo_label(txt_path, img_w, img_h, id_to_label, label_to_id=None): """Parse a YOLO-format .txt label file into AnyLabeling shape dicts. - Each line in the file is expected to follow the YOLO detection - format:: + Supports two YOLO formats — detection (bounding box) and segmentation + (polygon) — detected by the number of values on each line: + + * **Detection** (5 values):: + * **Segmentation** (7+ values, even coordinate count):: + + … + Coordinates are normalised (0-1). Returns a list of dicts ready to be passed to ``LabelingWidget.load_labels()``. @@ -75,21 +81,6 @@ def read_yolo_label(txt_path, img_w, img_h, id_to_label, label_to_id=None): continue class_id = int(parts[0]) - x_center = float(parts[1]) - y_center = float(parts[2]) - width = float(parts[3]) - height = float(parts[4]) - - # Normalised → absolute pixel coords - w_abs = width * img_w - h_abs = height * img_h - cx_abs = x_center * img_w - cy_abs = y_center * img_h - - x1 = cx_abs - w_abs / 2.0 - y1 = cy_abs - h_abs / 2.0 - x2 = cx_abs + w_abs / 2.0 - y2 = cy_abs + h_abs / 2.0 label = id_to_label.get(class_id) if label is None: @@ -98,15 +89,53 @@ def read_yolo_label(txt_path, img_w, img_h, id_to_label, label_to_id=None): if label_to_id is not None and label not in label_to_id: label_to_id[label] = class_id - shape = { - "label": label, - "text": "", - "points": [[x1, y1], [x2, y2]], - "shape_type": "rectangle", - "group_id": None, - "flags": {}, - "other_data": {}, - } + # Detection format (exactly 5 values) vs segmentation (7+ values) + if len(parts) == 5: + # YOLO detection: class_id x_center y_center width height + x_center = float(parts[1]) + y_center = float(parts[2]) + width = float(parts[3]) + height = float(parts[4]) + + # Normalised → absolute pixel coords + w_abs = width * img_w + h_abs = height * img_h + cx_abs = x_center * img_w + cy_abs = y_center * img_h + + x1 = cx_abs - w_abs / 2.0 + y1 = cy_abs - h_abs / 2.0 + x2 = cx_abs + w_abs / 2.0 + y2 = cy_abs + h_abs / 2.0 + + shape = { + "label": label, + "text": "", + "points": [[x1, y1], [x2, y2]], + "shape_type": "rectangle", + "group_id": None, + "flags": {}, + "other_data": {}, + } + elif len(parts) >= 7 and (len(parts) - 1) % 2 == 0: + # YOLO segmentation: class_id x1 y1 x2 y2 … xn yn + coords = [float(v) for v in parts[1:]] + points = [ + [coords[i] * img_w, coords[i + 1] * img_h] + for i in range(0, len(coords), 2) + ] + shape = { + "label": label, + "text": "", + "points": points, + "shape_type": "polygon", + "group_id": None, + "flags": {}, + "other_data": {}, + } + else: + continue + shapes.append(shape) return shapes @@ -118,6 +147,11 @@ def write_yolo_label(txt_path, shapes, img_w, img_h, label_to_id): *label_to_id* is mutated in-place: labels that are not already present receive the next unused integer class ID so that saving never silently drops annotations. + + Output format is *mixed-mode*: + * Rectangles → YOLO detection format (5 values) + * Polygons → YOLO segmentation format (variable values) + * Other shape types are skipped (not representable in YOLO format) """ used_ids = set(label_to_id.values()) next_id = 0 @@ -126,11 +160,8 @@ def write_yolo_label(txt_path, shapes, img_w, img_h, label_to_id): lines = [] for shape in shapes: - if shape["shape_type"] not in ("rectangle",): - continue + shape_type = shape["shape_type"] points = shape["points"] - if len(points) < 2: - continue label = shape["label"] @@ -143,18 +174,29 @@ def write_yolo_label(txt_path, shapes, img_w, img_h, label_to_id): class_id = label_to_id[label] - # Absolute pixel coords → normalised YOLO format - x1, y1 = points[0] - x2, y2 = points[1] - - x_center = ((x1 + x2) / 2.0) / img_w - y_center = ((y1 + y2) / 2.0) / img_h - width = abs(x2 - x1) / img_w - height = abs(y2 - y1) / img_h - - lines.append( - f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}" - ) + if shape_type == "rectangle" and len(points) >= 2: + # YOLO detection: class_id x_center y_center width height + x1, y1 = points[0] + x2, y2 = points[1] + + x_center = ((x1 + x2) / 2.0) / img_w + y_center = ((y1 + y2) / 2.0) / img_h + width = abs(x2 - x1) / img_w + height = abs(y2 - y1) / img_h + + lines.append( + f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}" + ) + + elif shape_type == "polygon" and len(points) >= 3: + # YOLO segmentation: class_id x1 y1 x2 y2 … xn yn + normalized = [] + for x, y in points: + normalized.append(f"{x / img_w:.6f}") + normalized.append(f"{y / img_h:.6f}") + lines.append(f"{class_id} {' '.join(normalized)}") + + # Circles, lines, linestrips, points — not representable in YOLO with open(txt_path, "w", encoding="utf-8") as f: f.write("\n".join(lines)) diff --git a/tests/test_yolo_io.py b/tests/test_yolo_io.py new file mode 100644 index 0000000..ac8029e --- /dev/null +++ b/tests/test_yolo_io.py @@ -0,0 +1,386 @@ +"""Tests for YOLO label I/O with polygon (segmentation) support.""" + +import os +import tempfile +import unittest + +import sys +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from anylabeling.views.labeling.yolo_io import ( + read_yolo_label, + write_yolo_label, +) + + +def _temp_path(suffix=".txt"): + fd, path = tempfile.mkstemp(suffix=suffix) + os.close(fd) + return path + + +class TestReadYoloLabel(unittest.TestCase): + """Tests for read_yolo_label() — detection and segmentation formats.""" + + def setUp(self): + self.img_w = 640 + self.img_h = 480 + self.id_to_label = {0: "person", 1: "car"} + self.tmp_path = _temp_path() + + def tearDown(self): + os.unlink(self.tmp_path) + + def _write(self, lines): + with open(self.tmp_path, "w", encoding="utf-8") as f: + f.write("\n".join(lines)) + + def test_read_detection_format(self): + """Detection format (5 values) → rectangle shape.""" + self._write(["0 0.5 0.5 0.4 0.3"]) + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), 1) + self.assertEqual(shapes[0]["shape_type"], "rectangle") + self.assertEqual(shapes[0]["label"], "person") + self.assertEqual(len(shapes[0]["points"]), 2) + + def test_read_segmentation_format_three_points(self): + """Segmentation format (7 values = 3 points) → polygon shape.""" + self._write(["0 0.1 0.1 0.5 0.1 0.3 0.4"]) + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), 1) + self.assertEqual(shapes[0]["shape_type"], "polygon") + self.assertEqual(shapes[0]["label"], "person") + self.assertEqual(len(shapes[0]["points"]), 3) + + def test_read_segmentation_format_five_points(self): + """Segmentation format (11 values = 5 points) → polygon shape.""" + self._write(["1 0.1 0.1 0.5 0.1 0.5 0.4 0.3 0.5 0.1 0.4"]) + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), 1) + self.assertEqual(shapes[0]["shape_type"], "polygon") + self.assertEqual(shapes[0]["label"], "car") + self.assertEqual(len(shapes[0]["points"]), 5) + + def test_read_mixed_formats(self): + """A single file with both detection and segmentation lines.""" + self._write([ + "0 0.5 0.5 0.4 0.3", + "1 0.1 0.1 0.5 0.1 0.3 0.4", + ]) + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), 2) + self.assertEqual(shapes[0]["shape_type"], "rectangle") + self.assertEqual(shapes[1]["shape_type"], "polygon") + + def test_read_normalised_coordinates(self): + """Check coordinate conversion from normalised → absolute pixels.""" + self._write(["0 0.5 0.5 0.4 0.3"]) + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + # x_center=0.5*640=320, width=0.4*640=256 → x1=320-128=192, x2=320+128=448 + # y_center=0.5*480=240, height=0.3*480=144 → y1=240-72=168, y2=240+72=312 + self.assertAlmostEqual(shapes[0]["points"][0][0], 192.0) + self.assertAlmostEqual(shapes[0]["points"][0][1], 168.0) + self.assertAlmostEqual(shapes[0]["points"][1][0], 448.0) + self.assertAlmostEqual(shapes[0]["points"][1][1], 312.0) + + def test_read_normalised_coordinates_polygon(self): + """Check polygon coordinate conversion.""" + self._write(["0 0.1 0.2 0.5 0.2 0.3 0.6"]) + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + # x: 0.1*640=64, 0.5*640=320, 0.3*640=192 + # y: 0.2*480=96, 0.2*480=96, 0.6*480=288 + expected = [[64.0, 96.0], [320.0, 96.0], [192.0, 288.0]] + for i, (x, y) in enumerate(expected): + with self.subTest(point=i): + self.assertAlmostEqual(shapes[0]["points"][i][0], x) + self.assertAlmostEqual(shapes[0]["points"][i][1], y) + + def test_read_populates_label_to_id(self): + """label_to_id is populated for new labels.""" + self._write(["3 0.5 0.5 0.4 0.3"]) + label_to_id = {} + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, + self.id_to_label, label_to_id=label_to_id, + ) + self.assertEqual(len(shapes), 1) + self.assertEqual(shapes[0]["label"], "class_3") + self.assertIn("class_3", label_to_id) + self.assertEqual(label_to_id["class_3"], 3) + + def test_read_unknown_class_id(self): + """An unseen class_id gets a fallback label.""" + self._write(["99 0.5 0.5 0.4 0.3"]) + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(shapes[0]["label"], "class_99") + + def test_read_skips_short_lines(self): + """Lines with < 5 values are skipped.""" + self._write([ + "0 0.5 0.5 0.4 0.3", + "1 0.1 0.2", + "2 0.5 0.5 0.4", + ]) + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), 1) + + def test_read_skips_odd_coord_count(self): + """Segmentation-like line with odd coord count is skipped.""" + self._write(["0 0.1 0.2 0.3 0.4 0.5"]) + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), 0) + + def test_read_skips_comments_and_blanks(self): + """Comments and blank lines are ignored.""" + self._write([ + "# this is a comment", + "", + "0 0.5 0.5 0.4 0.3", + " ", + "# another comment", + ]) + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), 1) + + def test_read_nonexistent_file(self): + """Missing file returns empty list.""" + shapes = read_yolo_label( + "/nonexistent/path.txt", self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(shapes, []) + + +class TestWriteYoloLabel(unittest.TestCase): + """Tests for write_yolo_label() — mixed-mode output.""" + + def setUp(self): + self.img_w = 640 + self.img_h = 480 + self.tmp_path = _temp_path() + + def tearDown(self): + os.unlink(self.tmp_path) + + def _read_lines(self): + with open(self.tmp_path, "r", encoding="utf-8") as f: + return [line.strip() for line in f if line.strip()] + + def test_write_rectangle(self): + """Rectangle shapes → detection format (5 values).""" + shapes = [{ + "label": "person", "shape_type": "rectangle", + "points": [[100, 80], [500, 320]], + "text": "", "group_id": None, "flags": {}, + }] + label_to_id = {"person": 0} + write_yolo_label(self.tmp_path, shapes, self.img_w, self.img_h, label_to_id) + lines = self._read_lines() + self.assertEqual(len(lines), 1) + parts = lines[0].split() + self.assertEqual(len(parts), 5) + self.assertEqual(parts[0], "0") + + def test_write_polygon(self): + """Polygon shapes → segmentation format (variable values).""" + shapes = [{ + "label": "person", "shape_type": "polygon", + "points": [[100, 80], [500, 80], [300, 320]], + "text": "", "group_id": None, "flags": {}, + }] + label_to_id = {"person": 0} + write_yolo_label(self.tmp_path, shapes, self.img_w, self.img_h, label_to_id) + lines = self._read_lines() + self.assertEqual(len(lines), 1) + parts = lines[0].split() + self.assertEqual(len(parts), 7) # class_id + 6 coords = 7 + self.assertEqual(parts[0], "0") + + def test_write_mixed_shapes(self): + """Mixed rectangle + polygon → mixed output.""" + shapes = [ + { + "label": "person", "shape_type": "rectangle", + "points": [[100, 80], [500, 320]], + "text": "", "group_id": None, "flags": {}, + }, + { + "label": "car", "shape_type": "polygon", + "points": [[50, 60], [200, 60], [200, 180], [50, 180]], + "text": "", "group_id": None, "flags": {}, + }, + ] + label_to_id = {"person": 0, "car": 1} + write_yolo_label(self.tmp_path, shapes, self.img_w, self.img_h, label_to_id) + lines = self._read_lines() + self.assertEqual(len(lines), 2) + parts0 = lines[0].split() + parts1 = lines[1].split() + self.assertEqual(len(parts0), 5) + self.assertEqual(len(parts1), 9) # class_id + 8 coords = 9 + self.assertEqual(parts0[0], "0") + self.assertEqual(parts1[0], "1") + + def test_write_auto_assigns_new_ids(self): + """Unseen labels get the next unused class ID.""" + shapes = [{ + "label": "dog", "shape_type": "rectangle", + "points": [[100, 80], [500, 320]], + "text": "", "group_id": None, "flags": {}, + }] + label_to_id = {"person": 0, "car": 1} + write_yolo_label(self.tmp_path, shapes, self.img_w, self.img_h, label_to_id) + self.assertEqual(label_to_id["dog"], 2) + + def test_write_skips_circle_shape(self): + """Non-rectangle/non-polygon shapes are skipped.""" + shapes = [ + { + "label": "person", "shape_type": "rectangle", + "points": [[100, 80], [500, 320]], + "text": "", "group_id": None, "flags": {}, + }, + { + "label": "dot", "shape_type": "point", + "points": [[200, 150]], + "text": "", "group_id": None, "flags": {}, + }, + ] + label_to_id = {"person": 0} + write_yolo_label(self.tmp_path, shapes, self.img_w, self.img_h, label_to_id) + lines = self._read_lines() + self.assertEqual(len(lines), 1) + self.assertNotIn("dot", " ".join(lines)) + + def test_write_empty_shapes(self): + """Empty shapes list produces empty (newline-free) file.""" + write_yolo_label(self.tmp_path, [], self.img_w, self.img_h, {}) + with open(self.tmp_path, "r", encoding="utf-8") as f: + content = f.read() + self.assertEqual(content, "") + + +class TestYoloRoundTrip(unittest.TestCase): + """End-to-end round-trip: read → write → read yields same data.""" + + def setUp(self): + self.img_w = 640 + self.img_h = 480 + self.id_to_label = {0: "person", 1: "car", 2: "dog"} + self.tmp_path = _temp_path() + + def tearDown(self): + os.unlink(self.tmp_path) + + def test_roundtrip_detection_only(self): + """Detection-format file round-trips with rectangle shapes.""" + original_lines = [ + "0 0.5 0.5 0.4 0.3", + "1 0.2 0.3 0.1 0.2", + ] + with open(self.tmp_path, "w", encoding="utf-8") as f: + f.write("\n".join(original_lines) + "\n") + + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + label_to_id = {"person": 0, "car": 1} + write_yolo_label( + self.tmp_path, shapes, self.img_w, self.img_h, label_to_id + ) + + shapes2 = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), len(shapes2)) + for s1, s2 in zip(shapes, shapes2): + self.assertEqual(s1["shape_type"], s2["shape_type"]) + self.assertEqual(s1["label"], s2["label"]) + for p1, p2 in zip(s1["points"], s2["points"]): + self.assertAlmostEqual(p1[0], p2[0], places=5) + self.assertAlmostEqual(p1[1], p2[1], places=5) + + def test_roundtrip_segmentation_only(self): + """Segmentation-format file round-trips with polygon shapes.""" + original_lines = [ + "0 0.1 0.1 0.5 0.1 0.3 0.4", + "2 0.2 0.2 0.6 0.2 0.6 0.6 0.2 0.6", + ] + with open(self.tmp_path, "w", encoding="utf-8") as f: + f.write("\n".join(original_lines) + "\n") + + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), 2) + label_to_id = {"person": 0, "dog": 2} + write_yolo_label( + self.tmp_path, shapes, self.img_w, self.img_h, label_to_id + ) + + shapes2 = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), len(shapes2)) + for i, (s1, s2) in enumerate(zip(shapes, shapes2)): + with self.subTest(shape=i): + self.assertEqual(s1["shape_type"], s2["shape_type"]) + self.assertEqual(s1["label"], s2["label"]) + self.assertEqual(len(s1["points"]), len(s2["points"])) + for j, (p1, p2) in enumerate(zip(s1["points"], s2["points"])): + with self.subTest(point=j): + self.assertAlmostEqual(p1[0], p2[0], places=5) + self.assertAlmostEqual(p1[1], p2[1], places=5) + + def test_roundtrip_mixed(self): + """Mixed detection+segmentation file round-trips correctly.""" + original_lines = [ + "0 0.5 0.5 0.4 0.3", + "1 0.1 0.1 0.5 0.1 0.3 0.4", + ] + with open(self.tmp_path, "w", encoding="utf-8") as f: + f.write("\n".join(original_lines) + "\n") + + shapes = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), 2) + self.assertEqual(shapes[0]["shape_type"], "rectangle") + self.assertEqual(shapes[1]["shape_type"], "polygon") + + label_to_id = {"person": 0, "car": 1} + write_yolo_label( + self.tmp_path, shapes, self.img_w, self.img_h, label_to_id + ) + + shapes2 = read_yolo_label( + self.tmp_path, self.img_w, self.img_h, self.id_to_label + ) + self.assertEqual(len(shapes), len(shapes2)) + for s1, s2 in zip(shapes, shapes2): + self.assertEqual(s1["shape_type"], s2["shape_type"]) + self.assertEqual(len(s1["points"]), len(s2["points"])) + + +if __name__ == "__main__": + unittest.main() From 206eb61f9304639677e697110936e9042605f55a Mon Sep 17 00:00:00 2001 From: Leon Kraim Date: Thu, 9 Jul 2026 23:56:40 +0200 Subject: [PATCH 4/6] Fix YOLO data.yaml not being updated and label dialog missing class names --- anylabeling/views/labeling/label_widget.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/anylabeling/views/labeling/label_widget.py b/anylabeling/views/labeling/label_widget.py index 4601c25..8d9a946 100644 --- a/anylabeling/views/labeling/label_widget.py +++ b/anylabeling/views/labeling/label_widget.py @@ -1417,8 +1417,6 @@ def reset_state(self): self.label_file = None self.other_data = {} self._yolo_label_path = None - self._yolo_config_path = None - self._yolo_config_type = None self.canvas.reset_state() def current_item(self): @@ -2228,6 +2226,17 @@ def load_file(self, filename=None): # noqa: C901 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(): @@ -2826,6 +2835,9 @@ def import_image_folder(self, dirpath, pattern=None, load=True): ) self._yolo_config_path = config_path self._yolo_config_type = config_type + # Populate the label dialog with YOLO class names + for name in self._yolo_label_to_id: + self.label_dialog.add_label_history(name) self.last_open_dir = dirpath self.filename = None From 1c14ce3ad993727436d91508452482ce419663f8 Mon Sep 17 00:00:00 2001 From: Leon Kraim Date: Fri, 10 Jul 2026 00:05:15 +0200 Subject: [PATCH 5/6] Fix in-memory _yolo_id_to_label not synced after update_dataset_config When save_labels() writes data.yaml with a new class via update_dataset_config(), the in-memory _yolo_id_to_label was never updated. On the next image switch + return, read_yolo_label() would look up the new class_id in the stale _yolo_id_to_label, find None, and fall back to class_N. --- anylabeling/views/labeling/label_widget.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/anylabeling/views/labeling/label_widget.py b/anylabeling/views/labeling/label_widget.py index 8d9a946..603e250 100644 --- a/anylabeling/views/labeling/label_widget.py +++ b/anylabeling/views/labeling/label_widget.py @@ -1880,6 +1880,9 @@ def format_shape(s): 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", From 241fbc65f9e7cd8a41c69b4b5fd89edb9174693b Mon Sep 17 00:00:00 2001 From: Leon Kraim Date: Fri, 10 Jul 2026 00:52:47 +0200 Subject: [PATCH 6/6] Populate unique_label_list with YOLO class names on import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit During YOLO dataset import, class names were only added to the label dialog's autocomplete history (add_label_history), but not to the unique_label_list sidebar. This made the label popup always start blank on the first shape, and forced the user to type or scroll the dropdown every time — unlike the non-YOLO workflow where clicking a label in the sidebar pre-fills it for the next shape. --- anylabeling/views/labeling/label_widget.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/anylabeling/views/labeling/label_widget.py b/anylabeling/views/labeling/label_widget.py index 603e250..ad1272b 100644 --- a/anylabeling/views/labeling/label_widget.py +++ b/anylabeling/views/labeling/label_widget.py @@ -2838,9 +2838,16 @@ def import_image_folder(self, dirpath, pattern=None, load=True): ) self._yolo_config_path = config_path self._yolo_config_type = config_type - # Populate the label dialog with YOLO class names + # 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