app.py

import gradio as gr
from gradio_bbox_annotator import BBoxAnnotator
import json
import os
from pathlib import Path
from PIL import Image
from io import BytesIO
import tempfile
import shutil
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 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())
MAX_SIZE = [1024, 1024]  # Maximum width and height for resized images

class ImageProcessor:
    def __init__(self):
        # Create a persistent directory for resized images
        self.base_dir = os.path.join(tempfile.gettempdir(), "annotation_tool")
        self.resized_dir = os.path.join(self.base_dir, "resized_images")
        self._setup_directories()
        logger.info(f"Initialized ImageProcessor with directory: {self.base_dir}")

    def _setup_directories(self):
        """Create necessary directories if they don't exist"""
        os.makedirs(self.resized_dir, exist_ok=True)
        logger.info(f"Set up directories: {self.resized_dir}")

    def resize_image(self, image_path):
        """Resize image and save to persistent directory"""
        try:
            logger.info(f"Processing image: {image_path}")
            
            # Read original image
            with open(image_path, "rb") as f:
                img = Image.open(BytesIO(f.read()))
                img.thumbnail(MAX_SIZE, Image.Resampling.LANCZOS)
            
            # Create a unique filename for the resized image
            original_filename = os.path.basename(image_path)
            resized_filename = f"resized_{original_filename}"
            resized_path = os.path.join(self.resized_dir, resized_filename)
            
            # Save resized image
            img.save(resized_path)
            logger.info(f"Saved resized image to: {resized_path}")
            
            return resized_path
            
        except Exception as e:
            logger.error(f"Error processing image: {str(e)}")
            raise

class AnnotationManager:
    def __init__(self):
        self.annotations = {}
        self.image_processor = ImageProcessor()
    
    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"
        
        # 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"
            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)}"
            
            # 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 process_upload(self, image_path):
        """Process uploaded image"""
        if not isinstance(image_path, (str, bytes, os.PathLike)):
            logger.warning(f"Invalid image path type: {type(image_path)}")
            return None
            
        try:
            logger.info(f"Processing upload: {image_path}")
            resized_path = self.image_processor.resize_image(image_path)
            logger.info(f"Successfully processed upload: {resized_path}")
            return resized_path
        except Exception as e:
            logger.error(f"Error in process_upload: {str(e)}")
            return None

    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(), f"❌ Error: {error_msg}"
            
        image_path, annotations = bbox_data
        # Use original filename (remove 'resized_' prefix)
        filename = os.path.basename(image_path)
        if filename.startswith("resized_"):
            filename = filename[8:]
            
        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
        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
    
    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(f"""
        # Advertisement and Text Annotation Tool
        
        **Instructions:**
        1. Upload an image (will be automatically resized to max {MAX_SIZE[0]}x{MAX_SIZE[1]})
        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
        
        **Annotation Limits per Image:**
        - advertisement: Maximum 1 annotation
        - text: Maximum 2 annotations
        """)
        
        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
        status_msg = gr.Markdown(label="Status")
        
        # Event handlers
        def handle_image_upload(bbox_data):
            try:
                if not bbox_data or not isinstance(bbox_data, tuple):
                    return None, "No image uploaded"
                
                image_path, annotations = bbox_data
                if not image_path:
                    return None, "No image path provided"
                    
                logger.info(f"Handling upload for: {image_path}")
                resized_path = annotation_mgr.process_upload(image_path)
                
                if resized_path and os.path.exists(resized_path):
                    logger.info(f"Processed image path: {resized_path}")
                    # Return the resized path and keep any existing annotations
                    return (resized_path, annotations)
                else:
                    error_msg = "Failed to process image"
                    logger.error(error_msg)
                    return None, error_msg
                    
            except Exception as e:
                error_msg = f"Error in upload handler: {str(e)}"
                logger.error(error_msg)
                return None, error_msg
        
        # Handle image upload and resizing
        bbox_input.upload(
            fn=handle_image_upload,
            inputs=[bbox_input],
            outputs=[bbox_input]
        )
        
        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()