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

# Load both models
model_pipeline_v1 = pipeline(task="image-classification", model="ppicazo/autotrain-ap-pass-fail-v1")
model_pipeline_v2 = pipeline(task="image-classification", model="ppicazo/allsky-stars-detected-v2")

def predict(image):
    # Resize the image to have width 1080 while keeping the aspect ratio
    width = 1080
    ratio = width / image.width
    height = int(image.height * ratio)
    resized_image = image.resize((width, height))
    
    # Perform predictions with both models
    predictions_v1 = model_pipeline_v1(resized_image)
    predictions_v2 = model_pipeline_v2(resized_image)
    
    # Format the results for each model
    results_v1 = {p["label"]: p["score"] for p in predictions_v1}
    results_v2 = {p["label"]: p["score"] for p in predictions_v2}
    
    # Return results as separate outputs
    return results_v1, results_v2

# Define the Gradio Interface
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload image"),
    outputs=[
        gr.Label(num_top_classes=5, label="Pass/Fail Model v1 Predictions"),
        gr.Label(num_top_classes=5, label="Stars Model v2 Predictions"),
    ],
    title="AP Classifier (Two Models)",
    allow_flagging="manual",
).launch()