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
51
52
53
54
55
56
57
| class FCN(tornado.web.RequestHandler):
def initialize(self, sess, pred, image):
self.sess = sess
self.pred = pred
self.image = image
def predict(self, mat):
preds = self.sess.run(self.pred, feed_dict={self.image: mat})
preds = postprocess(mat, preds)
return preds
def get_contour(self, mat):
nb_components, output, stats, centroids = cv2.connectedComponentsWithStats(mat, connectivity=8)
sizes = stats[1:, -1]
nb_components = nb_components - 1
min_size = 32
img2 = np.zeros((output.shape), dtype=np.uint8)
for i in range(0, nb_components):
if sizes[i] >= min_size:
img2[output == i + 1] = 255
mat2, contours, hierarchy = cv2.findContours(img2, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
return [np.squeeze(c).tolist() for c in contours]
def post(self):
fileinfo = self.request.files['image'][0]
fname = fileinfo['filename']
extn = os.path.splitext(fname)[-1]
cname = str(uuid.uuid4())+extn
with open(os.path.join(cur_path, __UPLOAD__, cname), 'w') as fh:
fh.write(fileinfo['body'])
filedata = fileinfo["body"]
mat = cv2.imdecode(np.asarray(bytearray(filedata), dtype=np.uint8), 1)
mat = cv2.resize(mat, (conf["width"], conf["height"]))
mat = np.array([mat])
res = self.predict(mat)
res = np.squeeze(res)
res = res.astype(np.uint8)
res = self.get_contour(res)
self.write(json.dumps({"result": res}))
if __name__ == "__main__":
os.environ["CUDA_VISIBLE_DEVICES"] = conf["CUDA"]
if sys.argv[1] == "train":
train()
elif sys.argv[1] == "test":
predict(conf["model_dir"])
else:
sess = tf.Session()
image = tf.placeholder(tf.float32, shape=[1, conf["height"], conf["width"], 3], name="input_image")
pred, _ = inference(image, 1)
saver = tf.train.Saver()
saver.restore(sess, conf["model_dir"])
logger.info("model ready")
application = tornado.web.Application([(r"/", FCN, dict(sess=sess, pred=pred, image=image))], debug=False)
application.listen(8080)
tornado.ioloop.IOLoop.instance().start()
|