Spaces:
Sleeping
Sleeping
File size: 4,714 Bytes
9a7dd44 |
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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import gradio as gr
from gradio_bbox_annotator import BBoxAnnotator
import json
import os
from pathlib import Path
# Define the allowed categories
CATEGORIES = ["advertisement", "text"]
class AnnotationManager:
def __init__(self):
self.annotations = {}
def add_annotation(self, bbox_data):
"""Add or update annotations for an image"""
if bbox_data and isinstance(bbox_data, tuple):
image_path, annotations = bbox_data
if isinstance(image_path, str) and annotations:
filename = os.path.basename(image_path)
formatted_annotations = []
for ann in annotations:
y1, y2, x1, x2, label = ann
# Ensure label is one of the allowed categories
if label not in CATEGORIES:
continue
formatted_annotations.append({
"annotation": [y1, y2, x1, x2],
"label": label
})
if formatted_annotations: # Only add if there are valid annotations
self.annotations[filename] = formatted_annotations
return self.get_json_annotations()
def get_json_annotations(self):
"""Get all annotations as formatted JSON string"""
return json.dumps(self.annotations, indent=2)
def clear_annotations(self):
"""Clear all annotations"""
self.annotations = {}
return ""
def create_interface():
annotation_mgr = AnnotationManager()
with gr.Blocks() as demo:
gr.Markdown("""
# Advertisement and Text Annotation Tool
**Instructions:**
1. Upload an image using the upload button in the annotator
2. Draw bounding boxes and select either 'advertisement' or 'text' as label
3. Click 'Save Annotations' to add to the collection
4. Repeat for all images
5. Click 'Download JSON' when finished
**Labels:**
- advertisement: Use for marking advertisement regions
- text: Use for marking text regions
""")
with gr.Row():
with gr.Column(scale=2):
bbox_input = BBoxAnnotator(
show_label=True,
label="Draw Bounding Boxes",
show_download_button=True,
interactive=True,
categories=CATEGORIES # Set the predefined categories
)
with gr.Column(scale=1):
json_output = gr.TextArea(
label="Combined Annotations JSON",
interactive=False,
lines=15
)
with gr.Row():
save_btn = gr.Button("Save Current Image Annotations", variant="primary")
clear_btn = gr.Button("Clear All Annotations", variant="secondary")
download_btn = gr.Button("Download Combined JSON", variant="primary")
# Add status message
status_msg = gr.Textbox(label="Status", interactive=False)
# Event handlers
def save_and_update(bbox_data):
if not bbox_data or not isinstance(bbox_data, tuple):
return annotation_mgr.get_json_annotations(), "No image or annotations to save"
image_path, annotations = bbox_data
if not annotations:
return annotation_mgr.get_json_annotations(), "No annotations drawn"
# Count annotations by type
ad_count = sum(1 for ann in annotations if ann[4] == "advertisement")
text_count = sum(1 for ann in annotations if ann[4] == "text")
json_str = annotation_mgr.add_annotation(bbox_data)
filename = os.path.basename(image_path) if isinstance(image_path, str) else "unknown"
return json_str, f"Saved for {filename}: {ad_count} advertisements, {text_count} text regions"
save_btn.click(
fn=save_and_update,
inputs=[bbox_input],
outputs=[json_output, status_msg]
)
clear_btn.click(
fn=lambda: (annotation_mgr.clear_annotations(), "All annotations cleared"),
inputs=[],
outputs=[json_output, status_msg]
)
download_btn.click(
fn=lambda x: {"text": x, "filename": "combined_annotations.json"},
inputs=[json_output],
outputs=[gr.File()]
)
return demo
if __name__ == "__main__":
demo = create_interface()
demo.launch() |