591 lines
22 KiB
Python
Executable File
591 lines
22 KiB
Python
Executable File
#!/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()
|