initial
This commit is contained in:
@@ -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.
|
||||
@@ -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()
|
||||
Executable
+590
@@ -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()
|
||||
Reference in New Issue
Block a user