File size: 5,347 Bytes
9a7dd44 fa27543 9a7dd44 8535ce3 fa27543 8535ce3 fa27543 8535ce3 9a7dd44 8535ce3 fa27543 8535ce3 fa27543 8535ce3 9a7dd44 fa27543 9a7dd44 fa27543 9a7dd44 8535ce3 9a7dd44 fa27543 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 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
import gradio as gr
from gradio_bbox_annotator import BBoxAnnotator
import json
import os
from pathlib import Path
# Define categories and their limits
CATEGORY_LIMITS = {
"advertisement": 1, # Maximum 1 advertisement annotation per image
"text": 2 # Maximum 2 text annotations per image
}
CATEGORIES = list(CATEGORY_LIMITS.keys())
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"
# Count annotations per category
category_counts = {cat: 0 for cat in CATEGORIES}
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)}"
# Count this annotation
category_counts[label] += 1
# Check category limits
for category, count in category_counts.items():
limit = CATEGORY_LIMITS[category]
if count > limit:
return False, f"Too many {category} annotations. Maximum allowed: {limit}"
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(), f"❌ Error: {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
counts = {cat: sum(1 for ann in annotations if ann[4] == cat) for cat in CATEGORIES}
counts_str = ", ".join(f"{count} {cat}" for cat, count in counts.items())
success_msg = f"✅ Successfully saved for {filename}: {counts_str}"
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 the appropriate label
3. Click 'Save Annotations' to add to the collection
4. Repeat for all images
5. Copy the combined JSON when finished
**Annotation Limits per Image:**
- advertisement: Maximum 1 annotation
- text: Maximum 2 annotations
""")
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() |