import gradio as gr
from transformers import pipeline
from PIL import Image
import os

# Lấy giá trị của biến môi trường   
model_name = os.getenv("Model_name")  # Lấy giá trị của biến môi trường  

# Khởi tạo pipeline  
pipe = pipeline("image-classification", model=model_name)  



# Hàm xử lý ảnh đầu vào
def detect(image):
    # Chạy mô hình trên ảnh
    results = pipe(image)
    
    # Chuyển kết quả thành dạng dễ đọc
    result_dict = {res["label"]: res["score"] for res in results}
    
    # Lấy xác suất của từng lớp (DeepFake và Real)
    deepfake_score = result_dict.get("Fake", 0.0)
    real_score = result_dict.get("Real", 0.0)
    
    return deepfake_score, real_score


custom_css = """
.button-gradient {
  background: linear-gradient(45deg, #ff416c, #ff4b2b, #ff9b00, #ff416c);
  background-size: 400% 400%;
  border: none;
  padding: 14px 28px;
  font-size: 16px;
  font-weight: bold;
  color: white;
  border-radius: 10px;
  cursor: pointer;
  transition: 0.3s ease-in-out;
  animation: gradientAnimation 2s infinite linear;
  box-shadow: 0 4px 10px rgba(255, 65, 108, 0.6);
}
@keyframes gradientAnimation {
  0% { background-position: 0% 50%; }
  25% { background-position: 50% 100%; }
  50% { background-position: 100% 50%; }
  75% { background-position: 50% 0%; }
  100% { background-position: 0% 50%; }
}
.button-gradient:hover {
  transform: scale(1.05);
  box-shadow: 0 6px 15px rgba(255, 75, 43, 0.8);
}
"""


# Giao diện Gradio
with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("# 🔍 DeepFake Detector\nUpload an image to check if it's DeepFake or Real.")
    
    with gr.Row():
        with gr.Column():
            image = gr.Image(type="pil", label="Upload Image")
            detect_button = gr.Button("Detect")
        
        with gr.Column():
            deepfake_label = gr.Label(label="DeepFake Probability")
            real_label = gr.Label(label="Real Probability")

    detect_button.click(detect, inputs=[image], outputs=[deepfake_label, real_label])

# Khởi chạy ứng dụng
demo.launch()