separate classification file for prediction only

classification/classification_predict.py

@@ -0,0 +1,89 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, models, transforms
+from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, random_split
+import os
+import matplotlib.pyplot as plt
+import random
+from PIL import Image
+import numpy as np
+import pandas as pd
+
+# THIS IS HARDCODED AND DEPENDS ON THE TRAINING. FOR THE MODEL THAT WE TRAINED AND IMPORT HERE, THESE ARE THE CLASSES THAT CAN BE PREDICTED
+class_names = list(['Annonaceae', 'Apocynaceae', 'Astrocarynum aculeatum', 'Attalea', 'Bactris', 'Bellucia', 'Bonomia', 'Byrsonima', 'Caryocar', 'Cecropia palmata', 'Cecropia ulei', 'Cordia', 'Couepia', 'Embaubas', 'Euterpe', 'Fabaceae', 'Genipa', 'Goupia', 'Guatteria', 'Humiria', 'Inga', 'Leopoldinia', 'Licania', 'Malvaceae', 'Meliaceae', 'Miconia', 'Myrcia', 'Ocotea', 'Oenocarpus', 'Pachira', 'Panopsis rubescens', 'Parkia', 'Passovia', 'Phoradendron', 'Pourouma', 'Remijia', 'Ruizterania', 'Salicaceae', 'Sapotaceae', 'Simarouba', 'Tachigali', 'Tapirira', 'Virola', 'Vochysia', 'Xylopia', 'cf. Swartzia'])
+num_classes = 46
+
+# Load the pre-trained ResNet50 model
+resnet50 = models.resnet50(weights='ResNet50_Weights.DEFAULT')
+
+# Freeze the parameters of the pre-trained model
+for param in resnet50.parameters():
+    param.requires_grad = False
+
+# Remove the final fully connected layer
+num_ftrs = resnet50.fc.in_features
+resnet50.fc = nn.Identity()  # Replace the final layer with an identity function to get the feature vectors
+
+# Define a custom neural network with one hidden layer and an output layer
+class CustomNet(nn.Module):
+    def __init__(self, num_ftrs, num_classes):
+        super(CustomNet, self).__init__()
+        self.resnet50 = resnet50
+        self.hidden = nn.Linear(num_ftrs, 512)
+        self.relu = nn.ReLU()
+        self.output = nn.Linear(512, num_classes)
+
+    def forward(self, x):
+        x = self.resnet50(x)  # Extract features using the pre-trained model
+        x = self.hidden(x)  # Pass through the hidden layer
+        x = self.relu(x)  # Apply ReLU activation
+        x = self.output(x)  # Output layer
+        return x
+
+
+# Function to load and preprocess the image
+def process_image(image_path):
+    data_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+    image = Image.open(image_path).convert('RGB')
+    image = data_transform(image)# data_transforms(image) # <-- data transforms uses all the random cropping as well
+    image = image.unsqueeze(0)  # Add batch dimension
+    return image
+
+
+# Function to predict the class of a single image
+def predict_single_image(image_path, model):
+    # Load the image and preprocess it
+    image = process_image(image_path)
+
+    # Load the model
+    model.eval()
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+
+    # Pass the image through the model
+    with torch.no_grad():
+        image = image.to(device)
+        outputs = model(image)
+        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
+
+    # Return the class names and their probabilities as a Pandas Series
+    return pd.Series(probabilities.cpu().numpy(), index=class_names).sort_values(ascending=False)
+
+def classify(img_path):
+    # Path to the single image
+    image_path = img_path
+
+    # Initialize your custom model
+    model = CustomNet(num_ftrs, num_classes)
+    # Load the trained model weights
+    model.load_state_dict(torch.load('./fine_tuned:plant_classifier.pth'))
+
+    # Predict the class probabilities
+    class_probabilities = predict_single_image(image_path, model)
+    return class_probabilities