Spaces:
Paused
Paused
| """ | |
| YOLO Exporter | |
| Exports image annotations to YOLO format: | |
| - One .txt file per image with lines: class_id cx cy w h (normalized 0-1) | |
| - classes.txt listing class names | |
| - data.yaml for Ultralytics compatibility | |
| """ | |
| import os | |
| import logging | |
| from typing import Optional, Tuple | |
| from .base import BaseExporter, ExportContext, ExportResult | |
| from .cv_utils import ( | |
| build_category_mapping, | |
| polygon_to_bbox, | |
| normalize_bbox, | |
| extract_image_annotations, | |
| get_image_dimensions, | |
| get_image_filename, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| class YOLOExporter(BaseExporter): | |
| format_name = "yolo" | |
| description = "YOLO format for object detection (Ultralytics compatible)" | |
| file_extensions = [".txt", ".yaml"] | |
| def can_export(self, context: ExportContext) -> Tuple[bool, str]: | |
| has_image_schema = any( | |
| s.get("annotation_type") == "image_annotation" | |
| for s in context.schemas | |
| ) | |
| if not has_image_schema: | |
| return False, "No image_annotation schema found in config" | |
| # Check that we can get image dimensions | |
| missing_dims = [] | |
| for ann in context.annotations: | |
| instance_id = ann.get("instance_id", "") | |
| item = context.items.get(instance_id, {}) | |
| img_anns = extract_image_annotations(ann) | |
| if img_anns: | |
| w, h = get_image_dimensions(item) | |
| if w <= 0 or h <= 0: | |
| missing_dims.append(instance_id) | |
| if missing_dims: | |
| return ( | |
| False, | |
| f"YOLO requires image dimensions. Missing for: " | |
| f"{', '.join(missing_dims[:5])}" | |
| f"{'...' if len(missing_dims) > 5 else ''}" | |
| ) | |
| return True, "" | |
| def export(self, context: ExportContext, output_path: str, | |
| options: Optional[dict] = None) -> ExportResult: | |
| options = options or {} | |
| warnings = [] | |
| files_written = [] | |
| category_map = build_category_mapping(context.annotations, context.schemas) | |
| labels_dir = os.path.join(output_path, "labels") | |
| os.makedirs(labels_dir, exist_ok=True) | |
| # Track which images have been written (handle multiple annotators) | |
| image_labels = {} # filename_stem -> list of label lines | |
| for ann in context.annotations: | |
| instance_id = ann.get("instance_id", "") | |
| item = context.items.get(instance_id, {}) | |
| img_anns = extract_image_annotations(ann) | |
| if not img_anns: | |
| continue | |
| img_w, img_h = get_image_dimensions(item) | |
| if img_w <= 0 or img_h <= 0: | |
| warnings.append(f"Skipping {instance_id}: no image dimensions") | |
| continue | |
| file_name = get_image_filename(item) or instance_id | |
| raw_stem = os.path.splitext(os.path.basename(file_name))[0] | |
| stem = "".join(c if c.isalnum() or c in "-_." else "_" for c in raw_stem) | |
| if stem not in image_labels: | |
| image_labels[stem] = [] | |
| for schema_name, objects in img_anns: | |
| for obj in objects: | |
| obj_type = obj.get("type", "") | |
| label = obj.get("label", "") | |
| if label not in category_map: | |
| warnings.append(f"Unknown label '{label}' in {instance_id}") | |
| continue | |
| class_id = category_map[label] | |
| if obj_type == "bbox": | |
| x = obj.get("x", 0) | |
| y = obj.get("y", 0) | |
| w = obj.get("width", 0) | |
| h = obj.get("height", 0) | |
| cx, cy, nw, nh = normalize_bbox(x, y, w, h, img_w, img_h) | |
| image_labels[stem].append( | |
| f"{class_id} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}" | |
| ) | |
| elif obj_type in ("polygon", "freeform"): | |
| points = obj.get("points", []) | |
| if not points: | |
| continue | |
| bx, by, bw, bh = polygon_to_bbox(points) | |
| cx, cy, nw, nh = normalize_bbox(bx, by, bw, bh, img_w, img_h) | |
| warnings.append( | |
| f"{obj_type} in {instance_id} converted to enclosing bbox" | |
| ) | |
| image_labels[stem].append( | |
| f"{class_id} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}" | |
| ) | |
| elif obj_type == "landmark": | |
| warnings.append( | |
| f"Landmark in {instance_id} skipped (not supported in YOLO)" | |
| ) | |
| else: | |
| warnings.append( | |
| f"Unknown type '{obj_type}' in {instance_id}" | |
| ) | |
| # Write label files | |
| for stem, lines in image_labels.items(): | |
| label_file = os.path.join(labels_dir, f"{stem}.txt") | |
| with open(label_file, "w") as f: | |
| f.write("\n".join(lines)) | |
| if lines: | |
| f.write("\n") | |
| files_written.append(label_file) | |
| # Write classes.txt | |
| sorted_labels = sorted(category_map.items(), key=lambda kv: kv[1]) | |
| classes_file = os.path.join(output_path, "classes.txt") | |
| with open(classes_file, "w") as f: | |
| for name, _ in sorted_labels: | |
| f.write(f"{name}\n") | |
| files_written.append(classes_file) | |
| # Write data.yaml for Ultralytics | |
| data_yaml = os.path.join(output_path, "data.yaml") | |
| with open(data_yaml, "w") as f: | |
| f.write(f"path: {output_path}\n") | |
| f.write("train: images/train\n") | |
| f.write("val: images/val\n") | |
| f.write(f"nc: {len(sorted_labels)}\n") | |
| f.write(f"names: [{', '.join(repr(n) for n, _ in sorted_labels)}]\n") | |
| files_written.append(data_yaml) | |
| return ExportResult( | |
| success=True, | |
| format_name=self.format_name, | |
| files_written=files_written, | |
| warnings=warnings, | |
| stats={ | |
| "num_images": len(image_labels), | |
| "num_annotations": sum(len(v) for v in image_labels.values()), | |
| "num_classes": len(sorted_labels), | |
| }, | |
| ) | |