# AdVision AI: Simple Ad Detector & Suggestion Tool for Hugging Face

import gradio as gr
from PIL import Image, ImageDraw
import torch
import requests
from io import BytesIO
import numpy as np
from ultralytics import YOLO

# Load YOLOv8 model (assumes pretrained for ad/banner detection)
# You can also use your custom model from Roboflow or YOLO
model = YOLO("yolov8n.pt")  # Change to custom ad-detection model if available

# Function to analyze ad positions and give suggestions
def analyze_ad_positions(detections, image_height):
    insights = []
    for box in detections:
        x1, y1, x2, y2 = box[:4]
        center_y = (y1 + y2) / 2

        if center_y < image_height * 0.33:
            pos = "Top"
            suggestion = "Good visibility ✅"
        elif center_y < image_height * 0.66:
            pos = "Middle"
            suggestion = "Moderate visibility ⚠️ Consider moving up"
        else:
            pos = "Bottom"
            suggestion = "Low visibility ❌ Move above the fold"

        insights.append(f"Ad at {pos} → {suggestion}")
    return insights

# Main prediction function
def detect_ads(image):
    img = image.convert("RGB")
    image_array = np.array(img)

    results = model(image_array)[0]  # First result
    detections = []
    draw = ImageDraw.Draw(img)

    for box in results.boxes.xyxy:
        x1, y1, x2, y2 = map(int, box.tolist())
        draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
        draw.text((x1, y1 - 10), "Ad", fill="red")
        detections.append((x1, y1, x2, y2))

    suggestions = analyze_ad_positions(detections, img.height)
    suggestions_text = "\n".join(suggestions) if suggestions else "No ads detected."

    return img, suggestions_text

# Gradio Interface
interface = gr.Interface(
    fn=detect_ads,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(type="pil"), gr.Textbox()],
    title="AdVision AI",
    description="Upload a webpage screenshot. Detects ad placements and gives visibility suggestions."
)

interface.launch()