dnth's picture
Update app.py
ec3e9d1 verified
import json
import gradio as gr
def process_metadata(file):
# Read the uploaded JSON file
with open(file.name, "r") as f:
data = json.load(f)
# First pass: collect all unique categories
unique_categories = set()
for item in data["media_items"]:
for metadata in item["metadata_items"]:
if metadata["type"] == "object_label":
unique_categories.add(metadata["properties"]["category_name"])
# Create category mapping dynamically
category_mapping = {cat: idx for idx, cat in enumerate(sorted(unique_categories))}
# Create a formatted string of the mapping
mapping_str = "\n".join([f"{cat}: {idx}" for cat, idx in category_mapping.items()])
results = []
# Process each media item
for item in data["media_items"]:
# Initialize lists for bboxes and categories
bboxes = []
categories = []
# Process each metadata item
for metadata in item["metadata_items"]:
if metadata["type"] == "object_label":
bbox = metadata["properties"]["bbox"]
bboxes.append(bbox)
category = category_mapping.get(
metadata["properties"]["category_name"], 0
)
categories.append(category)
# Create output format
output = {
"file_name": item["file_name"],
"objects": {"bbox": bboxes, "category": categories},
}
results.append(json.dumps(output))
# Return both results and mapping
return "\n".join(results), mapping_str
def save_text(text):
# Create a temporary file to save the output
temp_file = "metadata.jsonl"
with open(temp_file, "w") as f:
f.write(text)
return temp_file
# Create Gradio interface
with gr.Blocks(title="VL-HF Annotations Converter") as iface:
gr.Markdown("# VL-HF Annotations Converter")
gr.Markdown(
"Upload a JSON annotations file exported from Visual Layer and get a JSONL file to be used in Hugging Face AutoTrain."
)
with gr.Row():
input_file = gr.File(label="Upload JSON file exported from Visual Layer")
with gr.Row():
output_text = gr.Textbox(label="Converted JSONL output", lines=10, scale=3)
mapping_text = gr.Textbox(label="Category Mapping", lines=10)
with gr.Row():
download_btn = gr.DownloadButton(
label="Download JSONL", value="metadata.jsonl", interactive=True
)
# Connect the components
output = input_file.change(
fn=process_metadata, inputs=[input_file], outputs=[output_text, mapping_text]
).then( # Chain the save operation immediately after processing
fn=save_text, inputs=[output_text], outputs=[download_btn]
)
if __name__ == "__main__":
iface.launch()