File size: 1,069 Bytes
eecd883 |
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 35 36 |
from flask import Flask, render_template, request
from PIL import Image
from io import BytesIO
import base64
from predict import predict_potato
from model import model
import torch
model.load_state_dict(torch.load("models\\potato_model_statedict__f.pth", map_location=torch.device('cpu')))
app = Flask(__name__)
# Your predict_mask function here...
@app.route('/')
def home():
return render_template('index.html')
@app.route('/predict', methods=['POST'])
def predict():
# Get the image file from the request
file = request.files['file']
# Predict the mask
class_name, probability, image = predict_potato(file, model)
# Convert image to base64 format
buffered = BytesIO()
image.save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
# Pass the base64 encoded image to the frontend
return render_template('index.html', image=img_str, class_name=class_name, probability=probability)
if __name__ == '__main__':
app.run(debug=True)
|