projet-integrateur/IA/main.py
2022-01-19 21:49:22 +01:00

58 satır
1,3 KiB
Python

from flask import Flask # Pour créer le service
from flask import request # Pour faire des jolis "POST"
import requests # Pour "request" un server "autre"
import numpy as np
import json #return json : json.dumps(dico)
import torch
import os
from PIL import Image
#MS
app = Flask(__name__)
#app.config.from_object('config')
#TODO :Charger le model (ou le bon model si on complexifie)
model = torch.hub.load('.', 'custom', path='ai.pt', source='local') # load the AI from a local source
@app.route('/prediction', methods=['GET','POST'])
def prediction():
card = request.json["card"]
# Transform card in PIL Image
img = Image.fromarray(np.uint8(card))
prediction = model(img) # infere with a PIL image
#print("img7 predictions (pandas)")
print(prediction.pandas().xyxy[0]) # img1 predictions (pandas)
# Create a list of label of each image in the card
labels = []
res = prediction.pandas().xyxy[0].to_numpy()
for i in res:
# xmin ymin xmax ymax confidence class name
labels.append(i[-1]) #
print(labels)
nb = len(labels)
coords = [ [i[-7], i[-5], i[-6], i[-4]] for i in res ]
myJson = {
"nb" : nb,
"coord" : coords,
"label" : labels
}
return myJson
if( __name__ == "__main__"):
app.run(host="0.0.0.0", port=50001, debug=True)