Upload app.py

app.py

@@ -0,0 +1,63 @@
+#!/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])