Create app.py

app.py

@@ -0,0 +1,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()