Spaces:
Running
Running
File size: 1,095 Bytes
40ada70 35d6185 40ada70 |
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 |
import torch
import gradio as gr
from transformers import ConvNextForImageClassification, AutoImageProcessor
from PIL import Image
model = ConvNextForImageClassification.from_pretrained("facebook/convnext-base-224")
# Redefine classifier for 23 classes
model.classifier = torch.nn.Linear(in_features=1024, out_features=23)
# Load model configuration and weights manually
model.load_state_dict(torch.load("convnext_base_finetuned.pth", map_location="cpu")) # Load your finetuned weights
model.eval()
# Load the processor
processor = AutoImageProcessor.from_pretrained("facebook/convnext-base-224")
# Define a function to predict the class from an image
def predict(image):
# Preprocess the image
inputs = processor(images=image, return_tensors="pt")
# Perform inference
with torch.no_grad():
outputs = model(**inputs)
predicted_class = torch.argmax(outputs.logits, dim=1).item()
return predicted_class
# Create Gradio interface for user input
iface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs=gr.Textbox())
iface.launch()
|