Spaces:
Sleeping
Sleeping
File size: 3,508 Bytes
ae98750 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import gradio as gr
import torch
import numpy as np
import torch
import torchvision.models as models
from torch import nn
from albumentations import (
HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,
Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,
IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine, RandomResizedCrop,
IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose, Normalize, Cutout, CoarseDropout, ShiftScaleRotate, CenterCrop, Resize, Rotate,
ShiftScaleRotate, CenterCrop, Crop, Resize, Rotate, RandomShadow, RandomSizedBBoxSafeCrop,
ChannelShuffle, MotionBlur,Lambda,SmallestMaxSize
)
from albumentations.pytorch import ToTensorV2
# Your Albumentations transformations
test_transforms = Compose([
Resize(224, 224),
Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2(),
])
def load_cub200_classes():
"""
This function loads the classes from the classes.txt file and returns a dictionary
"""
with open("classes.txt", encoding="utf-8") as f:
classes = f.read().splitlines()
# convert classes to dictionary separating the lines by the first space
classes = {int(line.split(" ")[0]) : line.split(" ")[1] for line in classes}
# return the classes dictionary
return classes
def load_model():
"""
This function loads the trained model and returns it
"""
model = models.resnet50(pretrained=True)
# Freeze the initial layers up to a certain point (e.g., layer 4)
freeze_layers = 4
for idx, (name, param) in enumerate(model.named_parameters()):
if idx < freeze_layers:
param.requires_grad = False
# Replace the final fully connected layer for the new task
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 200)
# Add dropout layers
model = nn.Sequential(
model,
nn.Dropout(0.8), # Adjust the dropout rate as needed
nn.Linear(200, 200) # Add additional fully connected layer
)
# Load the state dictionary from the file
state_dict = torch.load("resnet50_7511.pt",map_location=torch.device('cpu'))
# Load the state dictionary into the model object
model.load_state_dict(state_dict)
# # Load actual model
# model = torch.load("resnet18.pth", map_location=torch.device('cpu'))
# set the model to evaluation mode
model.eval()
# return the model
return model
def predict_image(image):
"""
This function takes an image as input and returns the class label
"""
# load the model
model = load_model()
# model.eval()
# load the classes
classes = load_cub200_classes()
# Apply Albumentations transformations
transformed_image = test_transforms(image=image)['image']
# Convert the image to a PyTorch tensor
tensor_image = transformed_image.unsqueeze(0)
# Make prediction
with torch.no_grad():
output = model(tensor_image)
# Assuming the output is a tensor representing class probabilities
probabilities = torch.nn.functional.softmax(output[0], dim=0).numpy()
# Get the class with the highest probability
predicted_class = np.argmax(probabilities)
# Return the class label
return "Predicted Class: " + classes[predicted_class+1]
# create a gradio interface
gr.Interface(fn=predict_image, inputs="image", outputs="text").launch()
|