import gradio as gr
from PIL import Image
from transformers import pipeline
import numpy as np
import cv2

# Muat model deteksi AI
classifier = pipeline("image-classification", model="umm-maybe/AI-image-detector")

def calculate_blur(image):
    gray = np.array(image.convert("L"))
    return cv2.Laplacian(gray, cv2.CV_64F).var()

def calculate_noise(image):
    gray = np.array(image.convert("L"), dtype=np.float32)
    return np.std(gray - np.mean(gray))

def high_freq_artifacts(image):
    gray = np.array(image.convert("L"), dtype=np.float32)
    f = np.fft.fft2(gray)
    fshift = np.fft.fftshift(f)
    magnitude_spectrum = np.abs(fshift)
    rows, cols = gray.shape
    crow, ccol = rows // 2 , cols // 2
    r = min(crow, ccol) // 4
    mask = np.zeros((rows, cols))
    mask[crow-r:crow+r, ccol-r:ccol+r] = 1
    low_energy = np.sum(magnitude_spectrum * mask)
    total_energy = np.sum(magnitude_spectrum)
    high_ratio = (total_energy - low_energy) / total_energy * 100
    return high_ratio

def detect_image(image: Image.Image):
    blur_score = calculate_blur(image)
    noise_score = calculate_noise(image)
    high_freq_score = high_freq_artifacts(image)

    # Prediksi menggunakan model deteksi AI
    result = classifier(image)
    ai_score = result[0]['score'] if result[0]['label'] == 'AI-generated' else 0

    # Gabungkan skor
    final_score = ai_score * 0.7 + (1 - blur_score / 1000) * 0.2 + (1 - noise_score / 255) * 0.1

    if final_score > 0.5:
        return "🤖 Gambar ini kemungkinan besar dihasilkan oleh AI."
    else:
        return "✅ Gambar ini kemungkinan besar asli."

with gr.Blocks() as demo:
    gr.Markdown("### Deteksi Gambar AI vs Asli")
    with gr.Row():
        img_input = gr.Image(type="pil", label="Unggah Gambar")
        output = gr.Textbox(label="Hasil Deteksi")
    btn = gr.Button("Deteksi")
    btn.click(detect_image, inputs=img_input, outputs=output)

demo.launch()