utils.py

import torch
import torchvision.transforms as transforms
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import os
import streamlit as st
import pandas as pd


# Load and preprocess the image
def preprocess_image(image_path, transform):
    image = Image.open(image_path).convert("RGB")
    return image, transform(image).unsqueeze(0)

def extract_label_from_path(image_path):
    # Extract the parent folder name as the label
    label_name = os.path.basename(os.path.dirname(image_path))
    return label_name

# Predict using the model
def predict(model, image_tensor, device):
    model.eval()
    with torch.no_grad():
        image_tensor = image_tensor.to(device)
        outputs = model(image_tensor)
        # converts the raw data from our model(logits) into a probability distribution across classes
        # uses sotftmax function
        probabilities = torch.nn.functional.softmax(outputs, dim=1)
        # converts probabilities tensor into a numpy array
    return probabilities.cpu().numpy().flatten()

# display results
def visualize_predictions(original_image, probabilities, class_names, true_label = None):
    fig, axarr = plt.subplots(1, 2, figsize=(14, 7))

    # Display image
    axarr[0].imshow(original_image)
    axarr[0].axis("off")
    # Add the true label to the image title
    if true_label is not None:
      axarr[0].set_title(f"True Label: {true_label}", fontsize=14)
    # Display predictions
    axarr[1].barh(class_names, probabilities)
    axarr[1].set_xlabel("Probability")
    axarr[1].set_title("Class Predictions")
    axarr[1].set_xlim(0, 1)

    plt.tight_layout()
    plt.show()


# defining prediction functions for streamlit frontend visual which are slightly different
def preprocess_image_st(image, transform):
    image = image.convert("RGB")  # Ensure it's in RGB format
    image_tensor = transform(image).unsqueeze(0)  # Apply transformations and add batch dimension
    return image_tensor
    
def visualize_predictions_st(probabilities, class_names, true_label=None):

    # Display true label if available
    if true_label is not None:
        st.write(f"True Label: {true_label}")

    # Display predictions as a bar chart
    prediction_data = {
        "Class Names": list(class_names.values()),
         "Probabilities": probabilities
    }
    df = pd.DataFrame(prediction_data)

    st.bar_chart(df.set_index("Class Names")["Probabilities"])

    # Display a detailed table of predictions
    st.write("Class Probabilities:")
    st.dataframe(df)