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

# 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: Minimal Test Case to Verify Model and Processor
try:
    print("Running a minimal test case with the model...")

    # Test Image URL
    url = "https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png"
    image = Image.open(requests.get(url, stream=True).raw)

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

    # Process the image
    test_inputs = processor(text=test_categories, images=image, return_tensors="pt", padding=True)
    print(f"Test inputs processed: {test_inputs}")

    # Perform inference
    test_outputs = model(**test_inputs)
    print(f"Test outputs: {test_outputs}")

    # Check probabilities
    test_logits = test_outputs.logits_per_image
    test_probs = test_logits.softmax(dim=1)
    print(f"Test probabilities: {test_probs}")

except Exception as e:
    print(f"Error during the minimal test case: {e}")
    raise RuntimeError(f"Test case failed: {e}")

# Step 3: Define the Inference Function
def classify_image(image):
    """
    Classify an image as 'safe' or 'unsafe' and return probabilities.

    Args:
        image (PIL.Image.Image): Uploaded image.
    
    Returns:
        str: Predicted category.
        dict: Probabilities for "safe" and "unsafe".
    """
    try:
        print("Starting image classification...")

        # Check if the image is valid
        if image is None:
            raise ValueError("No image provided. Please upload a valid image.")
        if not hasattr(image, "convert"):
            raise ValueError("Uploaded file is not a valid image format.")

        # Define main categories
        categories = ["safe", "unsafe"]
        print(f"Categories: {categories}")

        # Process the image
        print("Processing the image with the processor...")
        inputs = processor(text=categories, images=image, return_tensors="pt", padding=True)
        print(f"Processed inputs: {inputs}")

        # Perform inference
        print("Running model inference...")
        outputs = model(**inputs)
        print(f"Model outputs: {outputs}")

        # Calculate probabilities
        logits_per_image = outputs.logits_per_image
        probs = logits_per_image.softmax(dim=1)
        print(f"Probabilities: {probs}")

        # Extract probabilities for each category
        safe_prob = probs[0][0].item() * 100
        unsafe_prob = probs[0][1].item() * 100

        # Determine the predicted category
        predicted_category = "safe" if safe_prob > unsafe_prob else "unsafe"
        print(f"Predicted category: {predicted_category}")

        # Return the predicted category and probabilities
        return predicted_category, {"safe": f"{safe_prob:.2f}%", "unsafe": f"{unsafe_prob:.2f}%"}

    except Exception as e:
        print(f"Error during classification: {e}")
        return f"Error: {str(e)}", {}

# Step 4: Set Up Gradio Interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Textbox(label="Predicted Category"),  # Display the predicted category prominently
        gr.Label(label="Probabilities"),        # Display probabilities with a progress bar
    ],
    title="Content Safety Classification",
    description="Upload an image to classify it as 'safe' or 'unsafe' with corresponding probabilities.",
)

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