Spaces:
Sleeping
Sleeping
| 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() |