Update app.py

app.py

@@ -4,32 +4,63 @@ import json
 import os
 from pathlib import Path
 
-# Define the allowed categories
 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"""
-        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()
+        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"""
@@ -38,7 +69,7 @@ class AnnotationManager:
     def clear_annotations(self):
         """Clear all annotations"""
         self.annotations = {}
-        return ""
+        return "", "All annotations cleared"
 
 def create_interface():
     annotation_mgr = AnnotationManager()
@@ -52,7 +83,7 @@ def create_interface():
         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
+        5. Copy the combined JSON when finished
         
         **Labels:**
         - advertisement: Use for marking advertisement regions
@@ -66,59 +97,37 @@ def create_interface():
                     label="Draw Bounding Boxes",
                     show_download_button=True,
                     interactive=True,
-                    categories=CATEGORIES  # Set the predefined categories
+                    categories=CATEGORIES
                 )
             
             with gr.Column(scale=1):
                 json_output = gr.TextArea(
                     label="Combined Annotations JSON",
-                    interactive=False,
-                    lines=15
+                    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")
-            download_btn = gr.Button("Download Combined JSON", variant="primary")
 
-        # Add status message
-        status_msg = gr.Textbox(label="Status", interactive=False)
+        # Add status message with different styling based on message type
+        status_msg = gr.Markdown(label="Status")
         
         # 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,
+            fn=annotation_mgr.add_annotation,
             inputs=[bbox_input],
             outputs=[json_output, status_msg]
         )
         
         clear_btn.click(
-            fn=lambda: (annotation_mgr.clear_annotations(), "All annotations cleared"),
+            fn=annotation_mgr.clear_annotations,
             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__":