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()