import gradio as gr
from transformers import CLIPModel, CLIPProcessor
from PIL import Image

# Step 1: Load Fine-Tuned Model from Hugging Face Model Hub
model_name = "quadranttechnologies/retail-content-safety-clip-finetuned"

print("Initializing the application...")

try:
    print("Loading the model from Hugging Face Model Hub...")
    model = CLIPModel.from_pretrained(model_name, trust_remote_code=True)
    processor = CLIPProcessor.from_pretrained(model_name)
    print("Model and processor loaded successfully.")
except Exception as e:
    print(f"Error loading the model or processor: {e}")
    raise RuntimeError(f"Failed to load model: {e}")

# Step 2: Define the Inference Function
def classify_image(image):
    """
    Classify an image as 'safe' or 'unsafe' and return probabilities.
    """
    try:
        if image is None:
            raise ValueError("No image provided. Please upload a valid image.")

        # Define categories
        categories = ["safe", "unsafe"]

        # Process the image
        inputs = processor(text=categories, images=image, return_tensors="pt", padding=True)

        # Run inference
        outputs = model(**inputs)

        # Extract logits
        logits_per_image = outputs.logits_per_image  # Shape: [1, 2]
        print(f"Logits: {logits_per_image}")

        # Apply softmax to logits to get probabilities
        probs = logits_per_image.softmax(dim=1)  # Shape: [1, 2]
        print(f"Softmax probabilities: {probs}")

        # Extract probabilities for each category
        safe_prob = probs[0][0].item()  # Extract 'safe' probability
        unsafe_prob = probs[0][1].item()  # Extract 'unsafe' probability
        print(f"Safe probability: {safe_prob}, Unsafe probability: {unsafe_prob}")

        # Normalize probabilities to ensure they sum to 100%
        total_prob = safe_prob + unsafe_prob
        print(f"Total probability before normalization: {total_prob}")
        safe_percentage = (safe_prob / total_prob) * 100
        unsafe_percentage = (unsafe_prob / total_prob) * 100

        # Ensure the sum is exactly 100%
        print(f"Normalized percentages: Safe={safe_percentage}%, Unsafe={unsafe_percentage}%")
        return {
            "safe": round(safe_percentage, 2),  # Rounded to 2 decimal places
            "unsafe": round(unsafe_percentage, 2)
        }

    except Exception as e:
        return {"Error": str(e)}




# Step 3: Set Up Gradio Interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=2),  # Use gr.Label to display probabilities with a bar-style visualization
    title="Content Safety Classification",
    description="Upload an image to classify it as 'safe' or 'unsafe' with corresponding probabilities.",
)

# Step 4: Launch Gradio Interface
if __name__ == "__main__":
    print("Launching the Gradio interface...")
    iface.launch()