kapu / app.py
devdata's picture
Create app.py
7f44fbd
raw
history blame
809 Bytes
from transformers import MobileNetV2FeatureExtractor, MobileNetV2ForImageClassification
from PIL import Image
# Replace with the path to the directory containing the model files
model_folder = "/path/to/your/model_folder"
# Load the feature extractor
feature_extractor = MobileNetV2FeatureExtractor.from_pretrained(model_folder)
# Load the model
model = MobileNetV2ForImageClassification.from_pretrained(model_folder)
# Load the image
image = Image.open("/path/to/your/image.jpg")
# Preprocess the image
inputs = feature_extractor(images=image, return_tensors="pt")
# Make predictions
outputs = model(**inputs)
logits = outputs.logits
# Model predicts one of the 1000 ImageNet classes
predicted_class_idx = logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])