kapu / app.py
devdata's picture
Update app.py
df7a0c9
raw
history blame
827 Bytes
from transformers import MobileNetV2FeatureExtractor, MobileNetV2ForImageClassification
from PIL import Image
# Replace with the path to the directory containing the model files
model_folder = "mobilnetV2_ftprecmodelshuf.weights.best.hdf5"
# 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])