# Import necessary libraries
import gradio as gr  # Gradio for creating web interfaces
from fastai.vision.all import *  # FastAI library for deep learning tasks

# Define a custom transformation class for converting images to grayscale
class GrayscaleTransform(Transform):
    """
    Custom transformation class to convert images to grayscale.
    This is used to ensure that the input images match the format
    used during model training.
    """
    def encodes(self, img: PILImage):
        """
        Convert the input image to grayscale.
        
        Args:
            img (PILImage): The input image in PIL format.
        
        Returns:
            PIL.Image: The grayscale version of the input image.
        """
        return img.convert("L")  # 'L' mode represents grayscale images

# Load the pre-trained model
learn = load_learner('clocker.pkl')
"""
load_learner function loads a saved FastAI learner object.
The 'clocker.pkl' file contains the trained model, including
its architecture, weights, and any necessary preprocessing steps.
"""

def classify_image(img):
    """
    Classify the input image using the loaded model.
    
    Args:
        img: The input image to be classified.
    
    Returns:
        dict: A dictionary containing the prediction probabilities for each class.
    """
    # Make a prediction using the loaded model
    pred, _, probs = learn.predict(img)
    # Return a dictionary with class probabilities
    return {
        "average woman": float(probs[0]),  # Probability for "average woman" class
        "transgender woman": float(probs[1])  # Probability for "transgender woman" class
    }

# Create the Gradio interface
iface = gr.Interface(
    fn=classify_image,  # The function to be called when the interface is used
    inputs=gr.Image(),  # Input component: an image upload widget
    outputs=gr.Label(num_top_classes=2),  # Output component: label with top 2 classes
    title="Transfem Clocker AI",  # Title of the web interface
    description="Upload an image of a woman and this will guess if she is trans.",  # Description of the interface
)
"""
gr.Interface creates a web interface for the model:
- fn: The function to be called when an image is uploaded
- inputs: Specifies that the input should be an image
- outputs: Displays the top 2 class probabilities as labels
- title and description: Provides context for users
"""

# Launches the interface
iface.launch()
"""
This starts the Gradio interface, making it accessible via a web browser.
it is my first ever AI web app!
"""