Update app.py

app.py

@@ -4,7 +4,12 @@ import json
 import os
 from pathlib import Path
 
-CATEGORIES = ["advertisement", "text"]
+# 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):
@@ -22,7 +27,8 @@ class AnnotationManager:
         if not annotations:
             return False, "No annotations drawn"
         
-        # Check each annotation
+        # 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"
@@ -36,13 +42,22 @@ class AnnotationManager:
             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(), error_msg
+            return self.get_json_annotations(), f"❌ Error: {error_msg}"
             
         image_path, annotations = bbox_data
         filename = os.path.basename(image_path)
@@ -56,9 +71,9 @@ class AnnotationManager:
         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"
+        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
     
@@ -69,7 +84,7 @@ class AnnotationManager:
     def clear_annotations(self):
         """Clear all annotations"""
         self.annotations = {}
-        return "", "All annotations cleared"
+        return "", "🗑️ All annotations cleared"
 
 def create_interface():
     annotation_mgr = AnnotationManager()
@@ -80,14 +95,14 @@ def create_interface():
         
         **Instructions:**
         1. Upload an image using the upload button in the annotator
-        2. Draw bounding boxes and select either 'advertisement' or 'text' as label
+        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
         
-        **Labels:**
-        - advertisement: Use for marking advertisement regions
-        - text: Use for marking text regions
+        **Annotation Limits per Image:**
+        - advertisement: Maximum 1 annotation
+        - text: Maximum 2 annotations
         """)
         
         with gr.Row():