Spaces:
Running
Running
# Load model directly | |
from transformers import AutoImageProcessor, AutoModelForImageClassification | |
from PIL import Image # Import the Image module | |
import torch # Import the torch module | |
import streamlit as st | |
st.title("Image Classification") | |
uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "png"]) | |
#image_path = "/content/cm5_2.jpg" # Store the path as a string | |
processor = AutoImageProcessor.from_pretrained("mateoluksenberg/dit-base-Classifier_CM05") | |
model = AutoModelForImageClassification.from_pretrained("mateoluksenberg/dit-base-Classifier_CM05") | |
image = Image.open(uploaded_file) # Load the image from the file path | |
inputs = processor(image, return_tensors="pt") # Pass the image object to the processor | |
with torch.no_grad(): # Use torch.no_grad() to disable gradient calculations | |
logits = model(**inputs).logits | |
# model predicts one of the 1000 ImageNet classes | |
predicted_label = logits.argmax(-1).item() | |
print(model.config.id2label[predicted_label]) |