detect_bubbles.py

#!/usr/bin/env python3
"""
Bubble Detection Module - Optimized
===================================

Uses YOLO (You Only Look Once) model to detect text bubbles in manga/comic images
with advanced filtering to prevent duplicates and overlapping detections.

Features:
- Confidence threshold filtering
- Non-Maximum Suppression (NMS) 
- Overlap filtering
- Size-based filtering

Author: MangaTranslator Team
License: MIT
"""

import torch.serialization
from ultralytics import YOLO
import numpy as np


def calculate_iou(box1, box2):
    """
    Calculate Intersection over Union (IoU) between two bounding boxes
    
    Args:
        box1, box2: [x1, y1, x2, y2, confidence, class_id]
        
    Returns:
        float: IoU value between 0 and 1
    """
    # Extract coordinates
    x1_1, y1_1, x2_1, y2_1 = box1[:4]
    x1_2, y1_2, x2_2, y2_2 = box2[:4]
    
    # Calculate intersection area
    x1_inter = max(x1_1, x1_2)
    y1_inter = max(y1_1, y1_2)
    x2_inter = min(x2_1, x2_2)
    y2_inter = min(y2_1, y2_2)
    
    if x2_inter <= x1_inter or y2_inter <= y1_inter:
        return 0.0
    
    intersection_area = (x2_inter - x1_inter) * (y2_inter - y1_inter)
    
    # Calculate union area
    area1 = (x2_1 - x1_1) * (y2_1 - y1_1)
    area2 = (x2_2 - x1_2) * (y2_2 - y1_2)
    union_area = area1 + area2 - intersection_area
    
    return intersection_area / union_area if union_area > 0 else 0.0


def filter_overlapping_bubbles(bubbles, iou_threshold=0.3):
    """
    Remove overlapping bubbles using IoU threshold
    
    Args:
        bubbles: List of bubble detections
        iou_threshold: IoU threshold for considering bubbles as overlapping
        
    Returns:
        list: Filtered list without overlapping bubbles
    """
    if len(bubbles) <= 1:
        return bubbles
    
    # Sort by confidence score (descending)
    bubbles = sorted(bubbles, key=lambda x: x[4], reverse=True)
    
    filtered_bubbles = []
    
    for i, bubble in enumerate(bubbles):
        # Check if this bubble overlaps significantly with any already accepted bubble
        is_duplicate = False
        
        for accepted_bubble in filtered_bubbles:
            iou = calculate_iou(bubble, accepted_bubble)
            if iou > iou_threshold:
                is_duplicate = True
                break
        
        if not is_duplicate:
            filtered_bubbles.append(bubble)
    
    return filtered_bubbles


def filter_by_size(bubbles, min_width=20, min_height=20, max_width_ratio=0.8, max_height_ratio=0.8, image_width=None, image_height=None):
    """
    Filter bubbles by size constraints
    
    Args:
        bubbles: List of bubble detections
        min_width, min_height: Minimum bubble dimensions
        max_width_ratio, max_height_ratio: Maximum size relative to image
        image_width, image_height: Image dimensions for ratio calculation
        
    Returns:
        list: Size-filtered bubbles
    """
    filtered_bubbles = []
    
    for bubble in bubbles:
        x1, y1, x2, y2 = bubble[:4]
        width = x2 - x1
        height = y2 - y1
        
        # Check minimum size
        if width < min_width or height < min_height:
            continue
        
        # Check maximum size ratio (if image dimensions provided)
        if image_width and image_height:
            width_ratio = width / image_width
            height_ratio = height / image_height
            
            if width_ratio > max_width_ratio or height_ratio > max_height_ratio:
                continue
        
        filtered_bubbles.append(bubble)
    
    return filtered_bubbles


def detect_bubbles(model_path, image_path, conf_threshold=0.25, iou_threshold=0.3, enable_nms=True):
    """
    Detect text bubbles in manga/comic images using YOLOv8 model with advanced filtering
    
    This function loads a pre-trained YOLO model and uses it to identify
    text bubble regions with duplicate removal and overlap filtering.
    
    Args:
        model_path (str): Path to the YOLO model file (.pt format)
        image_path (str): Path to the input image or PIL Image object
        conf_threshold (float): Confidence threshold for detections (0.0-1.0)
        iou_threshold (float): IoU threshold for overlap filtering (0.0-1.0)
        enable_nms (bool): Enable Non-Maximum Suppression in YOLO
        
    Returns:
        list: List of detected bubbles with format:
              [x1, y1, x2, y2, confidence_score, class_id]
              where (x1,y1) is top-left corner and (x2,y2) is bottom-right corner
              
    Note:
        - Lower conf_threshold = more detections (but more false positives)
        - Lower iou_threshold = more aggressive overlap removal
        - Results are filtered and sorted by confidence
    """
    # Load YOLO model with safe globals for security
    with torch.serialization.safe_globals([YOLO]):
        model = YOLO(model_path)

    # Configure model parameters
    model.overrides['conf'] = conf_threshold  # Confidence threshold
    model.overrides['iou'] = 0.7 if enable_nms else 1.0  # NMS IoU threshold
    model.overrides['agnostic_nms'] = False  # Class-agnostic NMS
    model.overrides['max_det'] = 100  # Maximum detections per image

    # Run detection on the image
    results = model(image_path)[0]

    # Extract bounding box data
    if results.boxes is None or len(results.boxes) == 0:
        print("⚠️ No bubbles detected")
        return []
    
    bubbles = results.boxes.data.tolist()
    original_count = len(bubbles)
    
    print(f"🔍 Initial detections: {original_count}")

    # Apply confidence filtering (additional safety check)
    bubbles = [bubble for bubble in bubbles if bubble[4] >= conf_threshold]
    conf_filtered_count = len(bubbles)
    
    if conf_filtered_count < original_count:
        print(f"🎯 After confidence filter ({conf_threshold}): {conf_filtered_count}")

    # Get image dimensions for size filtering
    image_height, image_width = None, None
    if hasattr(results, 'orig_shape'):
        image_height, image_width = results.orig_shape
    
    # Apply size filtering
    bubbles = filter_by_size(
        bubbles, 
        min_width=15, 
        min_height=15,
        max_width_ratio=0.9, 
        max_height_ratio=0.9,
        image_width=image_width,
        image_height=image_height
    )
    size_filtered_count = len(bubbles)
    
    if size_filtered_count < conf_filtered_count:
        print(f"📏 After size filter: {size_filtered_count}")

    # Apply overlap filtering (our custom IoU-based filter)
    bubbles = filter_overlapping_bubbles(bubbles, iou_threshold)
    final_count = len(bubbles)
    
    if final_count < size_filtered_count:
        print(f"🎯 After overlap filter (IoU {iou_threshold}): {final_count}")

    # Sort final results by Y coordinate (top to bottom) then by confidence
    bubbles = sorted(bubbles, key=lambda x: (x[1], -x[4]))
    
    # Print filtering summary
    removed_count = original_count - final_count
    if removed_count > 0:
        print(f"✨ Removed {removed_count} duplicate/overlapping bubbles ({removed_count/original_count*100:.1f}%)")
    
    return bubbles