app.py

import gradio as gr
import easyocr
import cv2
import numpy as np
from PIL import Image

# Initialize Deep Learning OCR once
reader = easyocr.Reader(['en'])

# Helper to calculate Intersection over Union (Overlap)
def calculate_iou(boxA, boxB):
    # box = [x_left, y_top, x_right, y_bottom]
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])
    interArea = max(0, xB - xA) * max(0, yB - yA)
    if interArea == 0: return 0.0
    boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
    boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
    iou = interArea / float(boxAArea + boxBArea - interArea)
    return iou

def validate_whatsapp_hybrid(image):
    if image is None:
        return {"valid": False, "reason": "No image provided"}

    image_np = np.array(image.convert('RGB'))
    height, width, _ = image_np.shape
    center_x = width / 2
    debug_log = []

    # --- PHASE 1: DEEP OCR & HEADER BARRIER ---
    try:
        # detail=1 returns [ [[tl,tr,br,bl]], text, conf ]
        ocr_results = reader.readtext(image_np, detail=1, paragraph=False)
    except Exception as e:
        return {"valid": False, "error": f"OCR Failed: {str(e)}"}

    header_keywords = ["block", "add to contacts", "report", "business account", 
                       "joined in", "not a contact", "encrypted", "end-to-end"]
    
    safe_chat_start_y = 0
    
    # Find the lowest header element
    for (bbox, text, conf) in ocr_results:
        if conf < 0.3: continue
        bottom_y = bbox[2][1]
        if any(k in text.lower() for k in header_keywords):
            if bottom_y > safe_chat_start_y:
                safe_chat_start_y = bottom_y

    # Add buffer or set default
    if safe_chat_start_y > 0:
        safe_chat_start_y += 30
        debug_log.append(f"Header barrier set at Y={int(safe_chat_start_y)}")
    else:
        safe_chat_start_y = height * 0.18
        debug_log.append(f"No header keywords. Using default barrier Y={int(safe_chat_start_y)}")

    # --- PHASE 2: OPENCV STICKER CANDIDATES ---
    # Convert to grayscale
    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
    # Blur slightly to remove image noise
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    # Adaptive Threshold (handles different background brightness)
    thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 25, 15)

    # Find contours (blobs)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    sticker_candidates = []
    min_sticker_size = width * 0.12 # Sticker must be at least 12% of screen width
    max_sticker_size = width * 0.60

    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        
        # Filter A: Must be below header
        if y < safe_chat_start_y: continue
        
        # Filter B: Size criteria (Stickers aren't tiny dots or massive screens)
        if w < min_sticker_size or h < min_sticker_size: continue
        if w > max_sticker_size or h > max_sticker_size: continue

        # Filter C: Aspect Ratio (Stickers are roughly square-ish, 0.6 to 1.6 ratio)
        aspect = w / float(h)
        if not (0.6 < aspect < 1.6): continue
        
        # This is a potential sticker or a text bubble
        sticker_candidates.append([x, y, x+w, y+h]) # Format: [left, top, right, bottom]

    # --- PHASE 3: DE-DUPLICATION (Separating Text vs. Stickers) ---
    final_objects = []

    # 3a. Process OCR Text Objects
    ocr_boxes_formatted = []
    system_words = ["today", "yesterday", "messages", "calls", "unread"]
    
    for (bbox, text, conf) in ocr_results:
        top_y = bbox[0][1]
        if top_y < safe_chat_start_y: continue
        if conf < 0.4 or len(text) < 2: continue
        
        # Center check for system dates
        text_center = bbox[0][0] + ((bbox[1][0] - bbox[0][0]) / 2)
        if abs(text_center - center_x) < (width * 0.12): continue
        if any(s in text.lower() for s in system_words): continue

        # Format: [left, top, right, bottom]
        l, t, r, b = bbox[0][0], bbox[0][1], bbox[2][0], bbox[2][1]
        ocr_boxes_formatted.append([l, t, r, b])
        final_objects.append({
            'type': 'text', 'text_content': text, 'left': l, 'top': t, 
            'center': l + (r-l)/2
        })

    # 3b. Process Sticker Candidates against Text Boxes
    for s_box in sticker_candidates:
        is_actually_text = False
        # Check if this sticker candidate overlaps significantly with any text box
        for t_box in ocr_boxes_formatted:
            # If IoU > 15%, assume this blob is just the container for the text
            if calculate_iou(s_box, t_box) > 0.15:
                is_actually_text = True
                break
        
        if not is_actually_text:
            # It's a blob with no text inside! Probably a sticker.
            l, t, r, b = s_box
            final_objects.append({
                'type': 'sticker', 'text_content': 'STICKER/IMAGE', 'left': l, 'top': t,
                'center': l + (r-l)/2
            })
            debug_log.append(f"Found Sticker Candidate at Y={t}, X={l}")

    # --- PHASE 4: UNIFIED DECISION LOOP ---
    # Sort everything by top Y position
    final_objects.sort(key=lambda x: x['top'])

    final_decision = None
    stranger_limit = width * 0.20
    me_limit = width * 0.40

    for obj in final_objects:
        # Double-check center alignment (redundant safety check)
        if abs(obj['center'] - center_x) < (width * 0.10):
             debug_log.append(f"Skipped Centered {obj['type']}")
             continue

        # --- ALIGNMENT CHECK ---
        if obj['left'] < stranger_limit:
            final_decision = True
            debug_log.append(f"✅ VALID (Stranger): Found Left-Aligned {obj['type']} at X={int(obj['left'])}")
            break
        elif obj['left'] > me_limit:
            final_decision = False
            debug_log.append(f"❌ INVALID (Me): Found Right-Aligned {obj['type']} at X={int(obj['left'])}")
            break

    # --- RETURN ---
    if final_decision is True:
        return {"valid": True, "reason": "First message/sticker is from stranger.", "debug": debug_log}
    elif final_decision is False:
        return {"valid": False, "reason": "First message/sticker is from you.", "debug": debug_log}
    else:
        return {"valid": False, "reason": "No valid messages or stickers detected.", "debug": debug_log}

iface = gr.Interface(
    fn=validate_whatsapp_hybrid,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(),
    title="WhatsApp Validator V13 (Hybrid Text + Sticker Detection)",
    description="Uses Deep Learning OCR for text and OpenCV Blob Detection for stickers."
)

if __name__ == "__main__":
    iface.launch()