Spaces:
Sleeping
Sleeping
File size: 4,741 Bytes
9a7dd44 8535ce3 9a7dd44 8535ce3 9a7dd44 8535ce3 9a7dd44 8535ce3 9a7dd44 8535ce3 9a7dd44 8535ce3 9a7dd44 8535ce3 9a7dd44 8535ce3 9a7dd44 8535ce3 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 127 128 129 130 131 132 133 134 135 |
import gradio as gr
from gradio_bbox_annotator import BBoxAnnotator
import json
import os
from pathlib import Path
CATEGORIES = ["advertisement", "text"]
class AnnotationManager:
def __init__(self):
self.annotations = {}
def validate_annotations(self, bbox_data):
"""Validate the annotation data and return (is_valid, error_message)"""
if not bbox_data or not isinstance(bbox_data, tuple):
return False, "No image or annotations provided"
image_path, annotations = bbox_data
if not isinstance(image_path, str):
return False, "Invalid image format"
if not annotations:
return False, "No annotations drawn"
# Check each annotation
for ann in annotations:
if len(ann) != 5:
return False, "Invalid annotation format"
y1, y2, x1, x2, label = ann
# Validate coordinates
if any(not isinstance(coord, (int, float)) for coord in [y1, y2, x1, x2]):
return False, "Invalid coordinate values"
# Validate label
if not label or label not in CATEGORIES:
return False, f"Invalid or missing label. Must be one of: {', '.join(CATEGORIES)}"
return True, ""
def add_annotation(self, bbox_data):
"""Add or update annotations for an image"""
is_valid, error_msg = self.validate_annotations(bbox_data)
if not is_valid:
return self.get_json_annotations(), error_msg
image_path, annotations = bbox_data
filename = os.path.basename(image_path)
formatted_annotations = []
for ann in annotations:
y1, y2, x1, x2, label = ann
formatted_annotations.append({
"annotation": [y1, y2, x1, x2],
"label": label
})
self.annotations[filename] = formatted_annotations
# 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")
success_msg = f"Successfully saved for {filename}: {ad_count} advertisements, {text_count} text regions"
return self.get_json_annotations(), success_msg
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 "", "All annotations cleared"
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. Copy the combined 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
)
with gr.Column(scale=1):
json_output = gr.TextArea(
label="Combined Annotations JSON",
interactive=True,
lines=15,
show_copy_button=True
)
with gr.Row():
save_btn = gr.Button("Save Current Image Annotations", variant="primary")
clear_btn = gr.Button("Clear All Annotations", variant="secondary")
# Add status message with different styling based on message type
status_msg = gr.Markdown(label="Status")
# Event handlers
save_btn.click(
fn=annotation_mgr.add_annotation,
inputs=[bbox_input],
outputs=[json_output, status_msg]
)
clear_btn.click(
fn=annotation_mgr.clear_annotations,
inputs=[],
outputs=[json_output, status_msg]
)
return demo
if __name__ == "__main__":
demo = create_interface()
demo.launch() |