Spaces:
Running
Running
File size: 3,913 Bytes
b7f710c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import json
import argparse
import os
def convert_json(input_path, output_path, reference_dir):
"""
Convert index_to_class JSON to name_to_image JSON with file paths, including all images in reference directory.
Args:
input_path (str): Path to input JSON file
output_path (str): Path to output JSON file
reference_dir (str): Directory containing reference images
"""
# Define valid image extensions
valid_extensions = ['jpg', 'png', 'jpeg']
# Read the input JSON file
try:
with open(input_path, 'r') as f:
index_to_class = json.load(f)
except FileNotFoundError:
raise FileNotFoundError(f"Input JSON file {input_path} does not exist")
except json.JSONDecodeError:
raise ValueError(f"Input JSON file {input_path} is invalid")
# Initialize dictionary for name to image path mapping
name_to_image = {}
# Process class names from input JSON
json_entries = 0
for key, value in index_to_class.items():
if not isinstance(value, str):
print(f"Warning: Skipping non-string class name {value} for key {key}")
continue
# Try each valid extension for the image file
image_path = ""
for ext in valid_extensions:
potential_path = os.path.join(reference_dir, f"{value}.{ext.lower()}")
if os.path.exists(potential_path):
image_path = potential_path
break
name_to_image[value] = image_path
json_entries += 1
# Scan reference directory for all images with valid extensions
dir_entries = 0
for filename in os.listdir(reference_dir):
# Check if file has a valid image extension
if any(filename.lower().endswith(f".{ext}") for ext in valid_extensions):
# Extract the base name (without extension)
base_name = os.path.splitext(filename)[0]
# Only add to dictionary if not already present (avoid overwriting JSON-derived entries)
if base_name not in name_to_image:
image_path = os.path.join(reference_dir, filename)
name_to_image[base_name] = image_path
dir_entries += 1
# Ensure output directory exists
os.makedirs(os.path.dirname(output_path), exist_ok=True)
# Write the output JSON file
try:
with open(output_path, 'w') as f:
json.dump(name_to_image, f, indent=4)
except Exception as e:
raise IOError(f"Failed to write output JSON file {output_path}: {str(e)}")
# Print notification
total_entries = len(name_to_image)
print(f"Successfully created reference image path dictionary with {total_entries} entries "
f"({json_entries} from input JSON, {dir_entries} from reference directory). "
f"Output saved to {output_path}")
if __name__ == "__main__":
# Set up argument parser
parser = argparse.ArgumentParser(description='Convert index_to_class JSON to name_to_image JSON with file paths, including all images in reference directory')
parser.add_argument('--input', default='./ckpts/index_to_class_mapping.json',
help='Input JSON file path')
parser.add_argument('--output', default='./data/reference_data/reference_image_data.json',
help='Output JSON file path')
parser.add_argument('--reference-dir', default='data/reference_data/images',
help='Directory containing reference images')
args = parser.parse_args()
# Ensure reference directory exists
if not os.path.exists(args.reference_dir):
raise FileNotFoundError(f"Reference directory {args.reference_dir} does not exist")
# Run the conversion
convert_json(args.input, args.output, args.reference_dir) |