"""
CS5330 Fall 2024
Lab 2 Texture Classification
Calvin Lo
"""

import gradio as gr
import numpy as np
import joblib 
from skimage.feature import local_binary_pattern, graycomatrix, graycoprops

IMAGE_SIZE_GLCM = 256
IMAGE_SIZE_LBP = 128
RADIUS = 1
N_POINTS = 8 * RADIUS
LBP_METHOD = "uniform"


def get_feature_vector(img, feature_type):

    # convert to grayscale
    img_gray = img.convert("L")

    # resize and get feature vector
    if feature_type == "GLCM":
        img_resized = image_resize(img_gray, IMAGE_SIZE_GLCM)
        feature_vector = compute_glcm_histogram_pil(img_resized)
    else:
        img_resized = image_resize(img_gray, IMAGE_SIZE_LBP)
        feature_vector = get_lbp_hist(np.array(img_resized),
                                      N_POINTS,
                                      RADIUS,
                                      LBP_METHOD)
    return [feature_vector]


def compute_glcm_histogram_pil(image, distances=[1],
                               angles=[0],
                               levels=8,
                               symmetric=True):

    # Convert the PIL image to a NumPy array
    image_np = np.array(image)

    # Quantize the grayscale image to the specified number of levels
    image_np = (image_np * (levels - 1) / 255).astype(np.uint8)

    # Compute the GLCM using skimage's graycomatrix function
    glcm = graycomatrix(image_np,
                        distances=distances,
                        angles=angles,
                        levels=levels,
                        symmetric=symmetric,
                        normed=True)

    # Extract GLCM properties
    homogeneity = graycoprops(glcm, 'homogeneity')[0, 0]
    correlation = graycoprops(glcm, 'correlation')[0, 0]

    # Create the feature vector
    feature_vector = np.array([homogeneity, correlation])

    return feature_vector


def get_lbp_hist(gray_image, n_points, radius, method):
    # Compute LBP for the image
    lbp = local_binary_pattern(gray_image, n_points, radius, method)
    lbp_hist, _ = np.histogram(lbp.ravel(), bins=np.arange(0, n_points + 3),
                               range=(0, n_points + 2))

    # Normalize the histogram
    lbp_hist = lbp_hist.astype("float")
    lbp_hist /= (lbp_hist.sum() + 1e-6)  # Normalized histogram
    return lbp_hist


def image_resize(img, n):
    # Crop the image to a square by finding the minimum dimension
    min_dimension = min(img.size)
    left = (img.width - min_dimension) / 2
    top = (img.height - min_dimension) / 2
    right = (img.width + min_dimension) / 2
    bottom = (img.height + min_dimension) / 2
    img = img.crop((left, top, right, bottom))
    img = img.resize((n, n))
    return img


def classify(img, feature_type):
    # To load the model later
    loaded_model = joblib.load(feature_type + '_model.joblib')

    feature_vector = get_feature_vector(img, feature_type)
    # Make predictions with the loaded model
    label = loaded_model.predict(feature_vector)

    return f"{label[0]}"


def main():
    # Gradio interface
    with gr.Blocks() as interface:
        gr.Markdown("## Image Texture Classifier")

        # Image upload input
        img_input = gr.Image(type="pil")

        # Dropdown for selecting classifier
        classifier_dropdown = gr.Dropdown(choices=["GLCM", "LBP"],
                                          label="Feature Vector Type",
                                          value="GLCM")

        # Button for classification
        classify_button = gr.Button("Classify Image")

        # Output label
        output_label = gr.Textbox(label="Predicted Texture")

        # Set up interaction
        classify_button.click(fn=classify,
                              inputs=[img_input, classifier_dropdown],
                              outputs=output_label)

    # Launch the interface
    interface.launch()


if __name__ == "__main__":
    main()