-
Notifications
You must be signed in to change notification settings - Fork 1
/
server.py
50 lines (39 loc) · 1.52 KB
/
server.py
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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import os
import json
import tensorflow as tf
import numpy as np
from PIL import Image
from flask import Flask, request, jsonify
from flask_socketio import SocketIO, emit
from werkzeug import secure_filename
incep_res_model = tf.keras.applications.inception_v3.InceptionV3()
incep_res_model._make_predict_function()
recyclable = json.load(open('/root/feedme/recyclable.json'))
app = Flask(__name__)
app.secret_key = b'feedme'
socketio = SocketIO(app)
def save_post_img(img):
# print(type(img))
img_path = 'uploads/' + secure_filename(img.filename)
img.save(img_path)
return img_path
@app.route('/predict', methods=['POST'])
def incep_res_predict():
img_path = save_post_img(request.files['image'])
im = tf.keras.preprocessing.image.load_img(img_path, target_size=(299, 299))
im_array = tf.keras.preprocessing.image.img_to_array(im)
im_batch = np.expand_dims(im_array, axis=0)
processed = tf.keras.applications.inception_v3.preprocess_input(im_batch.copy())
predictions = incep_res_model.predict(processed)
decoded = tf.keras.applications.inception_v3.decode_predictions(predictions)
result = {label: float(confidence) for _, label, confidence in decoded[0]}
data = sorted(result.keys(), key=lambda x: result[x])[-1]
final = [data, recyclable[data]]
socketio.emit('new', final, json=True)
return jsonify(final)
@socketio.on('connect')
def new_connection():
# Return all saved data
pass
if __name__ == '__main__':
socketio.run(app, host='0.0.0.0', port=8000, debug=True)