app.py

import gradio as gr
import cv2
import numpy as np
import onnxruntime as ort
import requests
from PIL import Image
import io

# Load logo image with alpha channel (PNG with transparency)
logo = cv2.imread('logo.png', cv2.IMREAD_UNCHANGED)

# 1. Load the ONNX model once (CPU)
session = ort.InferenceSession("license_plate_yolo8.onnx", providers=["CPUExecutionProvider"])
input_name = session.get_inputs()[0].name

# 2. Preprocessing logic
def preprocess(img, size=640):
    h0, w0 = img.shape[:2]
    r = size / max(h0, w0)
    new_w, new_h = int(w0 * r), int(h0 * r)
    resized = cv2.resize(img, (new_w, new_h))
    dw, dh = size - new_w, size - new_h
    dw, dh = dw / 2, dh / 2
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    padded = cv2.copyMakeBorder(resized, top, bottom, left, right,
                                cv2.BORDER_CONSTANT, value=(114,114,114))
    img_rgb = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB)
    img_norm = img_rgb.astype(np.float32) / 255.0
    img_trans = img_norm.transpose(2,0,1)[None, ...]
    return img_trans, r, (left, top), (h0, w0)

# 3. Postprocessing logic
def postprocess(preds, r, pad, orig_shape, conf_thres=0.25, iou_thres=0.45):
    pad_x, pad_y = pad
    orig_h, orig_w = orig_shape
    raw = preds[0].squeeze(0)
    dets = raw.T  # shape (N,5)

    boxes, scores = [], []
    for x_center, y_center, w, h, conf in dets:
        if conf < conf_thres:
            continue
        xc = (x_center - pad_x) / r
        yc = (y_center - pad_y) / r
        bw = w / r
        bh = h / r
        x1 = max(0, xc - bw/2)
        y1 = max(0, yc - bh/2)
        x2 = min(orig_w, xc + bw/2)
        y2 = min(orig_h, yc + bh/2)
        boxes.append([int(x1), int(y1), int(x2-x1), int(y2-y1)])
        scores.append(float(conf))

    idxs = cv2.dnn.NMSBoxes(boxes, scores, conf_thres, iou_thres)
    if len(idxs) > 0:
        idxs = idxs.flatten()
        boxes = [boxes[i] for i in idxs]
    else:
        boxes = []
    return boxes

# 4. Overlay logo in detected boxes
def overlay_logo(img, boxes):
    if logo is None:
        print("Logo image not found. Please ensure 'logo.png' is in the working directory.")
        return img
    for x, y, w, h in boxes:
        overlay = cv2.resize(logo, (w, h), interpolation=cv2.INTER_AREA)
        if overlay.shape[2] == 4:
            alpha = overlay[:, :, 3] / 255.0
            for c in range(3):
                img[y:y+h, x:x+w, c] = (alpha * overlay[:, :, c] +
                                        (1 - alpha) * img[y:y+h, x:x+w, c])
        else:
            img[y:y+h, x:x+w] = overlay
    return img

# 5. Helper to load image from URL
def load_image_from_url(url):
    resp = requests.get(url)
    img_pil = Image.open(io.BytesIO(resp.content)).convert("RGB")
    img = np.array(img_pil)[..., ::-1]
    return img

# 6. Combined pipeline: detect and overlay logo
def detect_and_overlay(uploaded_image, url):
    if uploaded_image is not None:
        img_bgr = uploaded_image[..., ::-1]
    elif url:
        img_bgr = load_image_from_url(url)
    else:
        return None

    inp, r, pad, orig = preprocess(img_bgr)
    preds = session.run(None, {input_name: inp})
    boxes = postprocess(preds, r, pad, orig)
    out_img = overlay_logo(img_bgr, boxes)
    return out_img[..., ::-1]

# 7. Gradio UI with preloaded examples

def main():
    image_input = gr.Image(type="numpy", label="Upload Image")
    # url_input = gr.Textbox(label="(OR) Image URL")

    # List your preloaded vehicle images here (place files in 'preload_images/' folder)
    examples = [
        ["test/1.jpg", ""],
        ["test/2.jpg", ""],
        ["test/3.jpg", ""],
        ["test/4.jpg", ""],
        ["test/5.jpg", ""],
        ["test/7.jpg", ""],
    ]

    iface = gr.Interface(
        fn=detect_and_overlay,
        inputs=[image_input],
        outputs=gr.Image(type="numpy", label="Output Image"),
        examples=examples,
        title="Fenix - License Plate logo",
        description="Upload a vehicle image or click a thumbnail.",
        allow_flagging="never"
    )
    iface.launch(share=True)

if __name__ == "__main__":
    main()