This commit is contained in:
2026-02-04 15:22:28 +07:00
commit 7e56948ece
3 changed files with 1044 additions and 0 deletions
+61
View File
@@ -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 01).
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.
+393
View File
@@ -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 (01 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 0180; 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()
+590
View File
@@ -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()