react / app.py
devdata's picture
Create app.py
c94c4b6 verified
raw
history blame
1.31 kB
from flask import Flask, request, jsonify
from fastai.vision.all import *
from huggingface_hub import from_pretrained_fastai
from PIL import Image
import io
app = Flask(__name__)
def classify_image(image_file):
# Load the trained model from Huggingface
learn = from_pretrained_fastai("devdatanalytics/commonbean")
# Open the image file
img = Image.open(image_file)
# Perform any necessary preprocessing on the image
# For example, resizing or normalization
img = img.resize((224, 224)) # Resize the image to match the model's input size
# Convert the image to a BytesIO object
img_bytes = io.BytesIO()
img.save(img_bytes, format='PNG')
img_bytes.seek(0)
# Perform the classification
pred_class, pred_idx, probs = learn.predict(img_bytes)
# Return the classification results
return f"Predicted Class: {pred_class}, Probability: {probs[pred_idx]:.2f}"
@app.route('/classify', methods=['POST'])
def classify():
if 'image' not in request.files:
return "No image file found", 400
image_file = request.files['image']
# Perform image classification
classification_results = classify_image(image_file)
return jsonify(classification_results)
if __name__ == '__main__':
app.run()