MoinulwithAI's picture
Create app.py
37b7c68 verified
raw
history blame
1.75 kB
import gradio as gr
import torch
from torchvision import transforms
from PIL import Image
import numpy as np
# Load the trained model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = YourModel() # Replace 'YourModel' with your actual model class
model.load_state_dict(torch.load('D:/Dataset/Cricket Bowl Grip/final_model.pth'))
model.to(device)
model.eval()
# Define the transformation to be applied to the input image
transform = transforms.Compose([
transforms.Resize((224, 224)), # Resize image to fit your model's input size
transforms.ToTensor(), # Convert image to tensor
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # Example normalize values for ImageNet
])
# Define a function for making predictions
def predict(image):
image = Image.fromarray(image) # Convert numpy array to PIL Image
image = transform(image).unsqueeze(0) # Apply transformations and add batch dimension
image = image.to(device)
with torch.no_grad():
outputs = model(image)
_, predicted = torch.max(outputs, 1)
# Map predicted label to class name
class_names = ['OUTSWING', 'STRAIGHT', 'BACK_OF_HAND', 'CARROM', 'CROSSSEAM',
'GOOGLY', 'INSWING', 'KNUCKLE', 'LEGSPIN', 'OFFSPIN']
predicted_label = class_names[predicted.item()]
return predicted_label
# Create the Gradio Interface
iface = gr.Interface(fn=predict,
inputs=gr.Image(type="numpy"), # Accepts image input
outputs=gr.Text(), # Output the predicted class label
live=True) # live=True enables prediction while image is being uploaded
# Launch the interface
iface.launch()