From 7e56948ece6a4fb58509809938785c4af5bbb6ad Mon Sep 17 00:00:00 2001 From: ariska Date: Wed, 4 Feb 2026 15:22:28 +0700 Subject: [PATCH] initial --- README.md | 61 ++++ augment_yolov9_dataset.py | 393 ++++++++++++++++++++++++ convert_yolo_to_labelme.py | 590 +++++++++++++++++++++++++++++++++++++ 3 files changed, 1044 insertions(+) create mode 100644 README.md create mode 100644 augment_yolov9_dataset.py create mode 100755 convert_yolo_to_labelme.py diff --git a/README.md b/README.md new file mode 100644 index 0000000..7171826 --- /dev/null +++ b/README.md @@ -0,0 +1,61 @@ +# YOLOv9 dataset augmentation + +Augment a YOLOv9-format dataset by creating new image and label files for horizontal flip, vertical flip, +10% hue, +30% contrast, and grayscale. Labels are updated correctly for flips; other augmentations copy labels unchanged. + +## Dataset layout + +Expected structure: + +``` +dataset/ +├── images/ # .jpg or .png +│ ├── img1.jpg +│ └── img2.jpg +└── labels/ # .txt, one per image, same base name + ├── img1.txt + └── img2.txt +``` + +YOLO label format: one line per object: `class_id x_center y_center width height` (normalized 0–1). + +If `images/` and `labels/` are not present, the script treats the given directory as containing both images and labels (flat layout). + +## Setup + +```bash +pip install -r requirements.txt +``` + +## Usage + +Augment in place (new files appear next to originals in `images/` and `labels/`): + +```bash +python augment_yolov9_dataset.py --dataset-dir ./dataset/train +``` + +Write augmented files to a separate directory (creates `train_aug/images/` and `train_aug/labels/`): + +```bash +python augment_yolov9_dataset.py --dataset-dir ./dataset/train --output-dir ./dataset/train_aug +``` + +Other options: + +- `--image-ext .png` — look for `.png` instead of `.jpg` +- `--suffixes hflip vflip` — run only horizontal and vertical flip (choices: `hflip`, `vflip`, `hue`, `contrast`, `gray`) +- `--dry-run` — print which files would be created without writing + +## Output naming + +For each image `img.jpg` with label `img.txt`, the script can create: + +| Augmentation | Image | Label | +|----------------|-----------------|-----------------| +| Horizontal flip| `img_hflip.jpg` | `img_hflip.txt` | +| Vertical flip | `img_vflip.jpg` | `img_vflip.txt` | +| Hue +10% | `img_hue.jpg` | `img_hue.txt` | +| Contrast +30% | `img_contrast.jpg` | `img_contrast.txt` | +| Grayscale | `img_gray.jpg` | `img_gray.txt` | + +Add these paths to your YOLOv9 data YAML or file lists to use the augmented set. diff --git a/augment_yolov9_dataset.py b/augment_yolov9_dataset.py new file mode 100644 index 0000000..7a5c359 --- /dev/null +++ b/augment_yolov9_dataset.py @@ -0,0 +1,393 @@ +#!/usr/bin/env python3 +""" +Augment a YOLOv9-format dataset by creating new image and label files for: +horizontal flip, vertical flip, +10% hue, +30% contrast, and grayscale. +""" + +from __future__ import annotations + +import argparse +import logging +import random +import time +from pathlib import Path + +import cv2 + +# Augmentation strength constants (tune as needed) +HUE_DELTA = 0.1 # 10% hue shift in [0, 1] scale +CONTRAST_FACTOR = 1.3 # 30% contrast increase + +# Suffix used for each augmentation type -> (suffix, applies to labels) +SUFFIX_HFLIP = "hflip" +SUFFIX_VFLIP = "vflip" +SUFFIX_HUE = "hue" +SUFFIX_CONTRAST = "contrast" +SUFFIX_GRAY = "gray" + +LOG = logging.getLogger(__name__) + +# Default image extensions to discover (case-insensitive) +DEFAULT_IMAGE_EXTS = (".jpg", ".jpeg", ".png") + + +def read_yolo_labels(path: Path) -> list[tuple[int, float, float, float, float]]: + """Read YOLO label file; return list of (class_id, x_center, y_center, width, height).""" + rows = [] + with path.open() as f: + for line in f: + line = line.strip() + if not line: + continue + parts = line.split() + if len(parts) != 5: + LOG.warning( + "Skipping malformed line in %s (expected 5 values): %s", + path, + line[:80], + ) + continue + class_id = int(parts[0]) + x_center = float(parts[1]) + y_center = float(parts[2]) + width = float(parts[3]) + height = float(parts[4]) + rows.append((class_id, x_center, y_center, width, height)) + return rows + + +def write_yolo_labels(path: Path, rows: list[tuple[int, float, float, float, float]]) -> None: + """Write YOLO label file in one-line-per-object format.""" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w") as f: + for class_id, x_center, y_center, width, height in rows: + f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n") + + +def flip_labels_horizontal( + rows: list[tuple[int, float, float, float, float]], +) -> list[tuple[int, float, float, float, float]]: + """Return new rows with x_center replaced by 1 - x_center.""" + return [(c, 1.0 - x, y, w, h) for c, x, y, w, h in rows] + + +def flip_labels_vertical( + rows: list[tuple[int, float, float, float, float]], +) -> list[tuple[int, float, float, float, float]]: + """Return new rows with y_center replaced by 1 - y_center.""" + return [(c, x, 1.0 - y, w, h) for c, x, y, w, h in rows] + + +def _load_image(path: Path): + """Load image as BGR; raise on failure.""" + img = cv2.imread(str(path)) + if img is None: + raise OSError(f"Failed to load image: {path}. Check path and format (e.g. .jpg, .png).") + return img + + +def _ensure_parent(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + + +def apply_horizontal_flip( + image_path: Path, + labels_path: Path, + out_image_path: Path, + out_labels_path: Path, + dry_run: bool = False, +) -> None: + """Flip image horizontally and transform labels (x_center -> 1 - x_center).""" + if dry_run: + LOG.info("Would create: %s, %s", out_image_path, out_labels_path) + return + img = _load_image(image_path) + flipped = cv2.flip(img, 1) + _ensure_parent(out_image_path) + if not cv2.imwrite(str(out_image_path), flipped): + raise OSError(f"Failed to write image: {out_image_path}. Check permissions and disk space.") + rows = read_yolo_labels(labels_path) + write_yolo_labels(out_labels_path, flip_labels_horizontal(rows)) + + +def apply_vertical_flip( + image_path: Path, + labels_path: Path, + out_image_path: Path, + out_labels_path: Path, + dry_run: bool = False, +) -> None: + """Flip image vertically and transform labels (y_center -> 1 - y_center).""" + if dry_run: + LOG.info("Would create: %s, %s", out_image_path, out_labels_path) + return + img = _load_image(image_path) + flipped = cv2.flip(img, 0) + _ensure_parent(out_image_path) + if not cv2.imwrite(str(out_image_path), flipped): + raise OSError(f"Failed to write image: {out_image_path}. Check permissions and disk space.") + rows = read_yolo_labels(labels_path) + write_yolo_labels(out_labels_path, flip_labels_vertical(rows)) + + +def apply_hue_shift( + image_path: Path, + labels_path: Path, + out_image_path: Path, + out_labels_path: Path, + delta: float = HUE_DELTA, + dry_run: bool = False, +) -> None: + """Shift hue by delta (0–1 scale); copy labels unchanged.""" + if dry_run: + LOG.info("Would create: %s, %s", out_image_path, out_labels_path) + return + img = _load_image(image_path) + hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype("float32") + h, s, v = cv2.split(hsv) + # OpenCV H is 0–180; treat delta as fraction of full circle + h = (h + delta * 180) % 180 + hsv = cv2.merge([h, s, v]).astype("uint8") + out = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) + _ensure_parent(out_image_path) + if not cv2.imwrite(str(out_image_path), out): + raise OSError(f"Failed to write image: {out_image_path}. Check permissions and disk space.") + rows = read_yolo_labels(labels_path) + write_yolo_labels(out_labels_path, rows) + + +def apply_contrast( + image_path: Path, + labels_path: Path, + out_image_path: Path, + out_labels_path: Path, + factor: float = CONTRAST_FACTOR, + dry_run: bool = False, +) -> None: + """Apply contrast: (pixel - mean) * factor + mean, clip to [0, 255]; copy labels.""" + if dry_run: + LOG.info("Would create: %s, %s", out_image_path, out_labels_path) + return + img = _load_image(image_path).astype("float32") + mean = img.mean() + out = (img - mean) * factor + mean + out = out.clip(0, 255).astype("uint8") + _ensure_parent(out_image_path) + if not cv2.imwrite(str(out_image_path), out): + raise OSError(f"Failed to write image: {out_image_path}. Check permissions and disk space.") + rows = read_yolo_labels(labels_path) + write_yolo_labels(out_labels_path, rows) + + +def apply_grayscale( + image_path: Path, + labels_path: Path, + out_image_path: Path, + out_labels_path: Path, + dry_run: bool = False, +) -> None: + """Convert to grayscale and broadcast to 3 channels; copy labels.""" + if dry_run: + LOG.info("Would create: %s, %s", out_image_path, out_labels_path) + return + img = _load_image(image_path) + gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + out = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR) + _ensure_parent(out_image_path) + if not cv2.imwrite(str(out_image_path), out): + raise OSError(f"Failed to write image: {out_image_path}. Check permissions and disk space.") + rows = read_yolo_labels(labels_path) + write_yolo_labels(out_labels_path, rows) + + +def discover_images_and_labels( + dataset_dir: Path, + image_ext: str, +) -> list[tuple[Path, Path]]: + """ + Find (image_path, label_path) pairs. + Prefer dataset_dir/images/ and dataset_dir/labels/; else use dataset_dir for both. + """ + images_dir = dataset_dir / "images" + labels_dir = dataset_dir / "labels" + if not images_dir.is_dir(): + images_dir = dataset_dir + labels_dir = dataset_dir + if not images_dir.is_dir(): + raise FileNotFoundError( + f"Dataset directory not found or has no 'images' subdir: {dataset_dir}. " + "Provide a path that contains an 'images' folder or is the folder with image files." + ) + if not labels_dir.is_dir(): + raise FileNotFoundError( + f"Labels directory not found: {labels_dir}. " + "Expected a 'labels' folder next to 'images', or the same folder for flat layout." + ) + pairs = [] + + raw = (image_ext or "").strip() + if raw.lower() in {"*", "any", "all", "auto"}: + allowed_exts = {e.lower() for e in DEFAULT_IMAGE_EXTS} + else: + parts = [p.strip() for p in raw.split(",") if p.strip()] + if not parts: + allowed_exts = {e.lower() for e in DEFAULT_IMAGE_EXTS} + else: + allowed_exts = {(p if p.startswith(".") else f".{p}").lower() for p in parts} + + for img_path in images_dir.iterdir(): + if not img_path.is_file(): + continue + if img_path.suffix.lower() not in allowed_exts: + continue + base = img_path.stem + label_path = labels_dir / f"{base}.txt" + if not label_path.is_file(): + LOG.warning("No label file for image %s, skipping: %s", img_path.name, label_path) + continue + pairs.append((img_path, label_path)) + return pairs + + +def run_augmentations( + dataset_dir: Path, + output_dir: Path | None, + image_ext: str, + enabled: set[str], + max_per_image: int, + dry_run: bool, +) -> None: + """Discover image/label pairs and apply up to max_per_image random augmentations per image.""" + pairs = discover_images_and_labels(dataset_dir, image_ext) + if not pairs: + LOG.warning("No image/label pairs found in %s with image-ext %s.", dataset_dir, image_ext) + return + enabled_list = list(enabled) + if not enabled_list: + LOG.warning("No augmentation types enabled.") + return + out_root = output_dir if output_dir is not None else dataset_dir + if output_dir is None: + out_images = dataset_dir / "images" if (dataset_dir / "images").is_dir() else dataset_dir + out_labels = dataset_dir / "labels" if (dataset_dir / "labels").is_dir() else dataset_dir + else: + out_images = out_root / "images" + out_labels = out_root / "labels" + total_images = len(pairs) + total_augmentations = 0 + start_time = time.perf_counter() + LOG.info("Starting augmentation: %d images, max %d per image.", total_images, max_per_image) + for idx, (img_path, label_path) in enumerate(pairs, start=1): + base = img_path.stem + ext = img_path.suffix + k = min(max_per_image, len(enabled_list)) + chosen = random.sample(enabled_list, k) + LOG.info( + "Processing image %d/%d: %s (%s)", + idx, + total_images, + img_path.name, + ", ".join(chosen), + ) + for suffix in chosen: + out_img = out_images / f"{base}_{suffix}{ext}" + out_lbl = out_labels / f"{base}_{suffix}.txt" + try: + if suffix == SUFFIX_HFLIP: + apply_horizontal_flip(img_path, label_path, out_img, out_lbl, dry_run=dry_run) + elif suffix == SUFFIX_VFLIP: + apply_vertical_flip(img_path, label_path, out_img, out_lbl, dry_run=dry_run) + elif suffix == SUFFIX_HUE: + apply_hue_shift(img_path, label_path, out_img, out_lbl, dry_run=dry_run) + elif suffix == SUFFIX_CONTRAST: + apply_contrast(img_path, label_path, out_img, out_lbl, dry_run=dry_run) + elif suffix == SUFFIX_GRAY: + apply_grayscale(img_path, label_path, out_img, out_lbl, dry_run=dry_run) + total_augmentations += 1 + except OSError as e: + LOG.error("Skipping %s %s: %s", suffix, img_path.name, e) + elapsed = time.perf_counter() - start_time + LOG.info( + "Completed: %d images, %d augmentations in %.1f s.", + total_images, + total_augmentations, + elapsed, + ) + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + parser = argparse.ArgumentParser( + description="Augment YOLOv9 dataset with flips, hue, contrast, and grayscale.", + ) + parser.add_argument( + "--dataset-dir", + type=Path, + required=True, + help="Root of the dataset (containing images/ and labels/ or flat image+label files).", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=None, + help="Where to write augmented files (default: same as dataset-dir).", + ) + parser.add_argument( + "--image-ext", + type=str, + default=",".join(DEFAULT_IMAGE_EXTS), + help=( + "Image extension(s) to look for. Provide a single ext (e.g. .jpg) or a comma-separated list " + "(e.g. .jpg,.jpeg,.png). Use 'all'/'any' to use the defaults." + ), + ) + parser.add_argument( + "--suffixes", + type=str, + nargs="+", + default=[SUFFIX_HFLIP, SUFFIX_VFLIP, SUFFIX_HUE, SUFFIX_CONTRAST, SUFFIX_GRAY], + choices=[SUFFIX_HFLIP, SUFFIX_VFLIP, SUFFIX_HUE, SUFFIX_CONTRAST, SUFFIX_GRAY], + help="Which augmentations can be applied (default: all).", + ) + parser.add_argument( + "--max-per-image", + type=int, + default=2, + metavar="N", + help="Maximum number of augmentation types to apply per image (default: 2).", + ) + parser.add_argument( + "--seed", + type=int, + default=None, + help="Random seed for reproducible augmentation selection.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Only print which files would be created.", + ) + args = parser.parse_args() + if args.max_per_image < 1: + LOG.error("--max-per-image must be at least 1.") + raise SystemExit(1) + if args.seed is not None: + random.seed(args.seed) + if not args.dataset_dir.is_dir(): + LOG.error( + "Dataset directory does not exist: %s. Create it and add images/ and labels/ (or image + label files).", + args.dataset_dir, + ) + raise SystemExit(1) + run_augmentations( + dataset_dir=args.dataset_dir, + output_dir=args.output_dir, + image_ext=args.image_ext, + enabled=set(args.suffixes), + max_per_image=args.max_per_image, + dry_run=args.dry_run, + ) + + +if __name__ == "__main__": + main() diff --git a/convert_yolo_to_labelme.py b/convert_yolo_to_labelme.py new file mode 100755 index 0000000..0467754 --- /dev/null +++ b/convert_yolo_to_labelme.py @@ -0,0 +1,590 @@ +#!/usr/bin/env python3 +""" +Script to convert YOLO txt label format to LabelMe JSON format. + +YOLO format: class_id x_center y_center width height (normalized 0.0-1.0) +LabelMe format: JSON with shapes containing rectangles with pixel coordinates +""" + +import os +import sys +import argparse +import json +import shutil +from pathlib import Path + +try: + from PIL import Image + HAS_PIL = True +except ImportError: + HAS_PIL = False + print("Warning: PIL/Pillow not installed. Image dimension detection required for conversion.") + print("Install with: pip install pillow") + + +def get_image_dimensions(image_path): + """Get image width and height.""" + if not HAS_PIL: + return None, None + try: + with Image.open(image_path) as img: + return img.size # Returns (width, height) + except Exception as e: + print(f"Warning: Could not read image {image_path}: {e}") + return None, None + + +def yolo_to_labelme_rectangle(x_center_norm, y_center_norm, width_norm, height_norm, + img_width, img_height): + """ + Convert YOLO normalized bounding box to LabelMe rectangle coordinates. + + Args: + x_center_norm, y_center_norm, width_norm, height_norm: Normalized coordinates (0.0-1.0) + img_width, img_height: Image dimensions in pixels + + Returns: + List of two points: [[x1, y1], [x2, y2]] for top-left and bottom-right corners + """ + # Denormalize center coordinates and dimensions + x_center = x_center_norm * img_width + y_center = y_center_norm * img_height + width = width_norm * img_width + height = height_norm * img_height + + # Calculate top-left and bottom-right corners + x1 = x_center - width / 2.0 + y1 = y_center - height / 2.0 + x2 = x_center + width / 2.0 + y2 = y_center + height / 2.0 + + # Ensure coordinates are within image bounds + x1 = max(0.0, min(img_width, x1)) + y1 = max(0.0, min(img_height, y1)) + x2 = max(0.0, min(img_width, x2)) + y2 = max(0.0, min(img_height, y2)) + + return [[float(x1), float(y1)], [float(x2), float(y2)]] + + +def is_normalized(value): + """Check if a coordinate value is normalized (0.0-1.0).""" + return 0.0 <= float(value) <= 1.0 + + +def find_image_file(txt_file, image_extensions=None): + """ + Find corresponding image file for a txt annotation file. + + Args: + txt_file: Path to txt annotation file + image_extensions: List of image extensions to try (default: ['.jpg', '.jpeg', '.png', '.bmp']) + + Returns: + Path to image file or None if not found + """ + if image_extensions is None: + image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff'] + + txt_file = Path(txt_file) + base_name = txt_file.stem + txt_dir = txt_file.parent + + # First, check if txt_file is in a 'labels' directory + # If so, look for corresponding 'images' directory + if txt_dir.name.lower() == 'labels': + # Try to find images directory at the same level + images_dir = txt_dir.parent / 'images' + if images_dir.exists(): + # Look for image in images directory + for ext in image_extensions: + potential_image = images_dir / f"{base_name}{ext}" + if potential_image.exists(): + return potential_image + # Try case variations + for ext in image_extensions: + for case_ext in [ext, ext.upper(), ext.capitalize()]: + potential_image = images_dir / f"{base_name}{case_ext}" + if potential_image.exists(): + return potential_image + + # Check in same directory as txt file + for ext in image_extensions: + potential_image = txt_dir / f"{base_name}{ext}" + if potential_image.exists(): + return potential_image + + # Check with case variations in same directory + for ext in image_extensions: + for case_ext in [ext, ext.upper(), ext.capitalize()]: + potential_image = txt_dir / f"{base_name}{case_ext}" + if potential_image.exists(): + return potential_image + + return None + + +def find_images_directory_for_labels(labels_dir): + """ + Find the corresponding images directory for a labels directory. + + Args: + labels_dir: Path to labels directory + + Returns: + Path to images directory or None if not found + """ + labels_dir = Path(labels_dir) + + # If the directory name is 'labels', look for 'images' at the same level + if labels_dir.name.lower() == 'labels': + images_dir = labels_dir.parent / 'images' + if images_dir.exists(): + return images_dir + + return None + + +def convert_yolo_to_labelme(txt_file, image_file=None, class_names=None, + image_extensions=None, include_image_data=False): + """ + Convert a single YOLO txt annotation file to LabelMe JSON format. + + Args: + txt_file: Path to YOLO txt annotation file + image_file: Path to corresponding image file (optional, will be searched if not provided) + class_names: Dictionary mapping class_id to class name (optional) + image_extensions: List of image extensions to search (default: ['.jpg', '.jpeg', '.png', '.bmp']) + include_image_data: Whether to include base64-encoded image data in JSON + + Returns: + Dictionary with LabelMe JSON structure + """ + txt_file = Path(txt_file) + + if not txt_file.exists(): + raise FileNotFoundError(f"Annotation file not found: {txt_file}") + + # Find image file if not provided + if image_file is None: + image_file = find_image_file(txt_file, image_extensions) + + if image_file is None: + raise FileNotFoundError( + f"Image file not found for {txt_file}. " + f"Please provide image_file or ensure image exists in same directory." + ) + + image_file = Path(image_file) + if not image_file.exists(): + raise FileNotFoundError(f"Image file not found: {image_file}") + + # Get image dimensions + img_width, img_height = get_image_dimensions(image_file) + if img_width is None or img_height is None: + raise ValueError( + f"Could not determine image dimensions for {image_file}. " + f"PIL/Pillow is required for this operation." + ) + + # Read YOLO annotations + shapes = [] + with open(txt_file, 'r') as f: + for line_num, line in enumerate(f, 1): + line = line.strip() + if not line: # Skip empty lines + continue + + parts = line.split() + if len(parts) < 5: + print(f"Warning: Invalid YOLO format in {txt_file} line {line_num}: {line}") + continue + + try: + class_id = int(parts[0]) + x_center = float(parts[1]) + y_center = float(parts[2]) + width = float(parts[3]) + height = float(parts[4]) + + # Check if coordinates are normalized + if not (is_normalized(x_center) and is_normalized(y_center) and + is_normalized(width) and is_normalized(height)): + print(f"Warning: Coordinates in {txt_file} line {line_num} may not be normalized. " + f"Assuming normalized format.") + + # Convert to LabelMe rectangle format + points = yolo_to_labelme_rectangle( + x_center, y_center, width, height, img_width, img_height + ) + + # Get class name + if class_names and class_id in class_names: + label = class_names[class_id] + else: + label = str(class_id) # Use class_id as label if no mapping provided + + # Create shape annotation + shape = { + "label": label, + "points": points, + "group_id": None, + "shape_type": "rectangle", + "flags": {} + } + shapes.append(shape) + + except (ValueError, IndexError) as e: + print(f"Warning: Could not parse line {line_num} in {txt_file}: {line} - {e}") + continue + + # Get image data if requested + image_data = None + if include_image_data: + try: + with open(image_file, 'rb') as f: + import base64 + image_data = base64.b64encode(f.read()).decode('utf-8') + except Exception as e: + print(f"Warning: Could not encode image data: {e}") + + # Create LabelMe JSON structure + labelme_json = { + "version": "5.0.1", + "flags": {}, + "shapes": shapes, + "imagePath": image_file.name, + "imageData": image_data, + "imageHeight": img_height, + "imageWidth": img_width + } + + return labelme_json + + +def convert_dataset(input_dir, output_dir=None, class_names_file=None, + image_extensions=None, include_image_data=False, + copy_images=False, recursive=False): + """ + Convert a directory of YOLO txt annotations to LabelMe JSON format. + + Args: + input_dir: Input directory containing txt files and images + output_dir: Output directory for LabelMe JSON files (optional, if None, JSON files are placed next to images) + class_names_file: Path to file with class names (one per line, optional) + image_extensions: List of image extensions to search + include_image_data: Whether to include base64-encoded image data + copy_images: Whether to copy images to output directory (only used if output_dir is specified) + recursive: Whether to process subdirectories recursively + + Returns: + Dictionary with conversion statistics + """ + input_dir = Path(input_dir) + + if not input_dir.exists(): + raise FileNotFoundError(f"Input directory not found: {input_dir}") + + # Load class names if provided + class_names = None + if class_names_file: + class_names_file = Path(class_names_file) + if class_names_file.exists(): + class_names = {} + with open(class_names_file, 'r') as f: + for idx, line in enumerate(f): + class_name = line.strip() + if class_name: + class_names[idx] = class_name + print(f"Loaded {len(class_names)} class names from {class_names_file}") + else: + print(f"Warning: Class names file not found: {class_names_file}") + + # Find all txt files (recursive or not) + if recursive: + txt_files = list(input_dir.rglob('*.txt')) + else: + txt_files = list(input_dir.glob('*.txt')) + + if not txt_files: + search_type = "recursively" if recursive else "in" + raise ValueError(f"No .txt files found {search_type} {input_dir}") + + stats = { + 'files_processed': 0, + 'total_annotations': 0, + 'errors': [] + } + + # Process each txt file + for txt_file in txt_files: + try: + # Find corresponding image + image_file = find_image_file(txt_file, image_extensions) + + if not image_file: + error_msg = f"Image file not found for {txt_file}" + stats['errors'].append(error_msg) + print(f"ERROR: {error_msg}") + continue + + # Convert to LabelMe format + labelme_json = convert_yolo_to_labelme( + txt_file, image_file, class_names, image_extensions, include_image_data + ) + + # Determine where to place the JSON file + if output_dir: + # If output_dir is specified, preserve relative path structure when recursive + output_dir = Path(output_dir) + if recursive: + # Preserve relative path from input_dir + relative_path = txt_file.relative_to(input_dir) + output_subdir = output_dir / relative_path.parent + output_subdir.mkdir(parents=True, exist_ok=True) + json_file = output_subdir / f"{image_file.stem}.json" + + # Copy image if requested, preserving directory structure + if copy_images: + output_image = output_subdir / image_file.name + if not output_image.exists(): + shutil.copy2(image_file, output_image) + else: + # Non-recursive: just use output_dir + output_dir.mkdir(parents=True, exist_ok=True) + json_file = output_dir / f"{image_file.stem}.json" + + # Copy image if requested + if copy_images: + output_image = output_dir / image_file.name + if not output_image.exists(): + shutil.copy2(image_file, output_image) + else: + # Check if txt_file is in a 'labels' directory + # If so, place JSON in corresponding 'images' directory + txt_dir = txt_file.parent + if txt_dir.name.lower() == 'labels': + images_dir = find_images_directory_for_labels(txt_dir) + if images_dir: + # Place JSON in images directory + json_file = images_dir / f"{image_file.stem}.json" + else: + # Fallback: place next to image file + json_file = image_file.parent / f"{image_file.stem}.json" + else: + # Otherwise, place JSON file next to the image file + json_file = image_file.parent / f"{image_file.stem}.json" + + json_file.parent.mkdir(parents=True, exist_ok=True) + with open(json_file, 'w') as f: + json.dump(labelme_json, f, indent=2) + + stats['files_processed'] += 1 + stats['total_annotations'] += len(labelme_json['shapes']) + + print(f"Processed: {txt_file} -> {json_file} ({len(labelme_json['shapes'])} annotations)") + + except Exception as e: + error_msg = f"Error processing {txt_file}: {str(e)}" + stats['errors'].append(error_msg) + print(f"ERROR: {error_msg}") + + return stats + + +def main(): + parser = argparse.ArgumentParser( + description='Convert YOLO txt label format to LabelMe JSON format', + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Convert single file (JSON placed in images folder if txt is in labels folder) + python convert_yolo_to_labelme.py train/labels/x.txt --image train/images/x.jpg + # Output: train/images/x.json + + # Convert directory (JSON files placed in images folders when txt files are in labels folders) + python convert_yolo_to_labelme.py --input-dir ./train/labels + # Converts train/labels/x.txt -> train/images/x.json + + # Convert directory recursively (processes all subdirectories) + python convert_yolo_to_labelme.py --input-dir ./dataset --recursive + # Converts train/labels/x.txt -> train/images/x.json + # Converts val/labels/y.txt -> val/images/y.json + + # Convert directory with custom output directory + python convert_yolo_to_labelme.py --input-dir ./labels --output-dir ./labelme_annotations + + # Convert recursively with custom output directory (preserves directory structure) + python convert_yolo_to_labelme.py --input-dir ./labels --output-dir ./labelme_annotations --recursive + + # Convert with class names file + python convert_yolo_to_labelme.py --input-dir ./labels --class-names classes.txt + + # Convert and include image data in JSON + python convert_yolo_to_labelme.py --input-dir ./labels --include-image-data + """ + ) + + parser.add_argument( + 'input', + nargs='?', + help='Input YOLO txt file (if converting single file)' + ) + + parser.add_argument( + '--image', + type=str, + help='Image file path (required for single file conversion)' + ) + + parser.add_argument( + '--output', '-o', + type=str, + help='Output JSON file path (for single file conversion). If not specified, JSON is placed in images folder when txt is in labels folder, otherwise next to image file.' + ) + + parser.add_argument( + '--input-dir', + type=str, + help='Input directory containing txt files and images (for batch conversion)' + ) + + parser.add_argument( + '--output-dir', + type=str, + default=None, + help='Output directory for LabelMe JSON files (for batch conversion). If not specified, JSON files are placed in the images folder when txt files are in a labels folder (e.g., train/labels/x.txt -> train/images/x.json), otherwise next to image files.' + ) + + parser.add_argument( + '--class-names', + type=str, + dest='class_names_file', + help='File with class names (one per line, line number = class_id)' + ) + + parser.add_argument( + '--image-extensions', + nargs='+', + default=['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff'], + help='Image file extensions to search for (default: .jpg .jpeg .png .bmp .tif .tiff)' + ) + + parser.add_argument( + '--include-image-data', + action='store_true', + help='Include base64-encoded image data in JSON (increases file size)' + ) + + parser.add_argument( + '--copy-images', + action='store_true', + help='Copy images to output directory (for batch conversion)' + ) + + parser.add_argument( + '--recursive', '-r', + action='store_true', + help='Process subdirectories recursively' + ) + + args = parser.parse_args() + + # Determine mode: single file or batch + if args.input: + # Single file mode + if not args.image: + parser.error("--image is required for single file conversion") + + # Load class names if provided + class_names = None + if args.class_names_file: + class_names_file = Path(args.class_names_file) + if class_names_file.exists(): + class_names = {} + with open(class_names_file, 'r') as f: + for idx, line in enumerate(f): + class_name = line.strip() + if class_name: + class_names[idx] = class_name + else: + print(f"Warning: Class names file not found: {class_names_file}") + + try: + labelme_json = convert_yolo_to_labelme( + args.input, + args.image, + class_names, + args.image_extensions, + args.include_image_data + ) + + # Determine output file path + if args.output: + output_file = Path(args.output) + else: + # Check if txt file is in a 'labels' directory + # If so, place JSON in corresponding 'images' directory + txt_file = Path(args.input) + txt_dir = txt_file.parent + image_file = Path(args.image) + + if txt_dir.name.lower() == 'labels': + images_dir = find_images_directory_for_labels(txt_dir) + if images_dir: + # Place JSON in images directory + output_file = images_dir / f"{image_file.stem}.json" + else: + # Fallback: place next to image file + output_file = image_file.parent / f"{image_file.stem}.json" + else: + # Place JSON file next to the image file + output_file = image_file.parent / f"{image_file.stem}.json" + + output_file.parent.mkdir(parents=True, exist_ok=True) + + with open(output_file, 'w') as f: + json.dump(labelme_json, f, indent=2) + + print(f"Successfully converted {args.input} to {output_file}") + print(f" Annotations: {len(labelme_json['shapes'])}") + print(f" Image: {labelme_json['imagePath']} ({labelme_json['imageWidth']}x{labelme_json['imageHeight']})") + + except Exception as e: + print(f"ERROR: {e}", file=sys.stderr) + sys.exit(1) + + elif args.input_dir: + # Batch mode - output_dir is optional + + try: + stats = convert_dataset( + args.input_dir, + args.output_dir, + args.class_names_file, + args.image_extensions, + args.include_image_data, + args.copy_images, + args.recursive + ) + + print("\n" + "="*50) + print("Conversion Summary:") + print(f" Files processed: {stats['files_processed']}") + print(f" Total annotations: {stats['total_annotations']}") + if stats['errors']: + print(f" Errors: {len(stats['errors'])}") + for error in stats['errors']: + print(f" - {error}") + print("="*50) + + except Exception as e: + print(f"ERROR: {e}", file=sys.stderr) + sys.exit(1) + + else: + parser.error("Either provide input file or --input-dir for batch conversion") + + +if __name__ == '__main__': + main()