#!/usr/bin/env python
# coding: utf-8
# Save this file as streamlit.py and run it using the command: streamlit run streamlit.py

import streamlit as st
import torch
from torch import nn
from PIL import Image
import numpy as np
import torchvision.transforms as transforms

# Custom model definition
class CustomModel(nn.Module):
    def __init__(self):
        super(CustomModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(32 * 224 * 224, 90)  # Adjust the dimensions and number of classes if necessary

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = x.view(x.size(0), -1)  # Flatten the tensor
        x = self.fc(x)
        return x

# Function to process the image
def process_image(img):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    img = transform(img)
    img = img.unsqueeze(0)  # Add batch dimension
    return img

st.title('Animal Classification')
st.write('Please choose an image so that the AI model can predict the type of animal.')
file = st.file_uploader('Pick an image', type=['jpg', 'jpeg', 'png'])

# Load animal names
with open("name of the animals.txt") as f:
    class_names = [x.strip() for x in f.readlines()]

if file is not None:
    img = Image.open(file)
    st.image(img, caption='The image: ')
    image = process_image(img)

    # Load the model
    model = CustomModel()
    model.load_state_dict(torch.load('model-CNN.pth', map_location=torch.device('cpu')))  # Load state dict
    model.eval()
    
    # Predict with the model
    with torch.no_grad():
        prediction = model(image)
    
    predicted_class = torch.argmax(prediction, dim=1).item()
    st.write('Probability Distribution')
    st.write(prediction.numpy())
    st.write("Prediction: ", class_names[predicted_class])