app.py

from __future__ import annotations

import base64
import json
import os
import smtplib
from datetime import datetime, timezone
from email.message import EmailMessage
from io import BytesIO

import gradio as gr
import requests
from PIL import Image

from aminatron_model import detect_image, format_detections, render_result_image, summarize_detections


SPACE_MODEL = os.getenv("AMINATRON_MODEL") or None
EMAIL_TO = os.getenv("EMAIL_TO", "mammadov.amin2000@gmail.com")
EMAIL_FROM = os.getenv("EMAIL_FROM", "Aminatron <onboarding@resend.dev>")
RESEND_API_KEY = os.getenv("RESEND_API_KEY")
SMTP_HOST = os.getenv("SMTP_HOST", "smtp.gmail.com")
SMTP_PORT = int(os.getenv("SMTP_PORT", "587"))
SMTP_USER = os.getenv("SMTP_USER")
SMTP_PASSWORD = os.getenv("SMTP_PASSWORD") or os.getenv("GMAIL_APP_PASSWORD")


def image_to_png_bytes(image: Image.Image) -> bytes:
    """Convert a PIL image to PNG bytes for email attachment."""

    buffer = BytesIO()
    image.convert("RGB").save(buffer, format="PNG")
    return buffer.getvalue()


def build_email_payload(
    original: Image.Image,
    annotated: Image.Image,
    detections: list[dict[str, object]],
    summary: dict[str, int],
) -> tuple[str, str, bytes, bytes]:
    """Build reusable email subject/body/attachments."""

    timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
    subject = f"Aminatron prediction review - {timestamp}"
    body = (
        "A new Aminatron Space prediction was processed.\n\n"
        f"Time: {timestamp}\n"
        f"Summary: {json.dumps(summary, ensure_ascii=False)}\n\n"
        f"Detections:\n{json.dumps(detections, ensure_ascii=False, indent=2)}\n"
    )
    return subject, body, image_to_png_bytes(original), image_to_png_bytes(annotated)


def send_review_email_resend(
    original: Image.Image,
    annotated: Image.Image,
    detections: list[dict[str, object]],
    summary: dict[str, int],
) -> str:
    """Send review email through Resend HTTPS API."""

    if not RESEND_API_KEY:
        return "Resend is not configured. Add RESEND_API_KEY to Space secrets."

    subject, body, original_png, result_png = build_email_payload(
        original,
        annotated,
        detections,
        summary,
    )
    response = requests.post(
        "https://api.resend.com/emails",
        headers={
            "Authorization": f"Bearer {RESEND_API_KEY}",
            "Content-Type": "application/json",
        },
        json={
            "from": EMAIL_FROM,
            "to": [EMAIL_TO],
            "subject": subject,
            "text": body,
            "attachments": [
                {
                    "filename": "original.png",
                    "content": base64.b64encode(original_png).decode("ascii"),
                },
                {
                    "filename": "aminatron_result.png",
                    "content": base64.b64encode(result_png).decode("ascii"),
                },
            ],
        },
        timeout=30,
    )
    if response.status_code >= 400:
        return f"Resend email failed: {response.status_code} {response.text}"

    return f"Email sent to {EMAIL_TO} through Resend."


def send_review_email(
    original: Image.Image,
    annotated: Image.Image,
    detections: list[dict[str, object]],
    summary: dict[str, int],
) -> str:
    """Send original and processed images to email."""

    if RESEND_API_KEY:
        try:
            return send_review_email_resend(original, annotated, detections, summary)
        except Exception as error:
            error_message = f"Resend email failed: {error}"
            print(error_message)
            return error_message

    if not SMTP_USER or not SMTP_PASSWORD:
        return (
            "Email is not configured. Add RESEND_API_KEY, or SMTP_USER and "
            "SMTP_PASSWORD/GMAIL_APP_PASSWORD."
        )

    subject, body, original_png, result_png = build_email_payload(
        original,
        annotated,
        detections,
        summary,
    )
    message = EmailMessage()
    message["Subject"] = subject
    message["From"] = SMTP_USER
    message["To"] = EMAIL_TO
    message.set_content(body)
    message.add_attachment(
        original_png,
        maintype="image",
        subtype="png",
        filename="original.png",
    )
    message.add_attachment(
        result_png,
        maintype="image",
        subtype="png",
        filename="aminatron_result.png",
    )

    try:
        smtp_class = smtplib.SMTP_SSL if SMTP_PORT == 465 else smtplib.SMTP
        with smtp_class(SMTP_HOST, SMTP_PORT, timeout=20) as smtp:
            if SMTP_PORT != 465:
                smtp.starttls()
            smtp.login(SMTP_USER, SMTP_PASSWORD)
            smtp.send_message(message)
        return f"Email sent to {EMAIL_TO}."
    except Exception as error:
        error_message = f"Email notification failed: {error}"
        print(error_message)
        return error_message


def predict(image: Image.Image, confidence: float, iou: float, image_size: int, max_detections: int):
    """Detect objects for the Gradio interface."""

    if image is None:
        return None, {}, [], None

    result = detect_image(
        image,
        model_path=SPACE_MODEL,
        confidence=confidence,
        iou=iou,
        image_size=image_size,
        max_detections=max_detections,
    )
    detections = format_detections(result)
    summary = summarize_detections(detections)
    annotated_image = render_result_image(result)
    rows = [
        [item["class"], f"{item['confidence'] * 100:.2f}%", item["box"]]
        for item in detections
    ]

    email_status = send_review_email(image, annotated_image, detections, summary)
    print(email_status)
    share_state = {
        "original": image,
        "annotated": annotated_image,
        "detections": detections,
        "summary": summary,
    }
    return annotated_image, summary, rows, share_state


def share_result(share_state):
    """Send the latest prediction result to email manually from the Share button."""

    if not share_state:
        print("No result to share yet. Run Detect first.")
        return

    email_status = send_review_email(
        share_state["original"],
        share_state["annotated"],
        share_state["detections"],
        share_state["summary"],
    )
    print(email_status)


with gr.Blocks(title="Aminatron") as demo:
    gr.Markdown(
        """
        # Aminatron

        Aminatron is a multi-object detection model. Upload one image and it can find
        several COCO objects at once: people, cats, dogs, birds, cows, cars, chairs and more.
        """
    )

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload image")
        image_output = gr.Image(type="pil", label="Detected objects")

    with gr.Row():
        confidence = gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Confidence")
        iou = gr.Slider(0.1, 0.9, value=0.45, step=0.05, label="IoU")
        image_size = gr.Dropdown([416, 512, 640, 768], value=640, label="Image size")
        max_detections = gr.Slider(1, 100, value=50, step=1, label="Max detections")

    share_state = gr.State(None)

    with gr.Row():
        detect_button = gr.Button("Detect")
        share_button = gr.Button("Share result")

    summary_output = gr.JSON(label="Summary")
    detections_output = gr.Dataframe(
        headers=["Class", "Confidence", "Box [x1, y1, x2, y2]"],
        label="Detections",
        interactive=False,
    )

    detect_button.click(
        fn=predict,
        inputs=[image_input, confidence, iou, image_size, max_detections],
        outputs=[image_output, summary_output, detections_output, share_state],
    )

    share_button.click(
        fn=share_result,
        inputs=[share_state],
        outputs=[],
    )


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