import sys

sys.path.append('./')

from yolo.net.yolo_tiny_net import YoloTinyNet

import tensorflow as tf

import cv2

import numpy as np

classes_name = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train","tvmonitor"]

def process_predicts(predicts):

return xmin, ymin, xmax, ymax, class_num

common_params = {'image_size': 448, 'num_classes': 20,

'batch_size':1}

net_params = {'cell_size': 7, 'boxes_per_cell':2, 'weight_decay': 0.0005}

net = YoloTinyNet(common_params, net_params, test=True)

image = tf.placeholder(tf.float32, (1, 448, 448, 3))

predicts = net.inference(image)

sess = tf.Session()

np_img = cv2.imread('cat.jpg')

resized_img = cv2.resize(np_img, (448, 448))

np_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB)

np_img = np_img.astype(np.float32)

np_img = np_img / 255.0 * 2 - 1

np_img = np.reshape(np_img, (1, 448, 448, 3))

saver = tf.train.Saver(net.trainable_collection)

saver.restore(sess, 'models/pretrain/yolo_tiny.ckpt')

np_predict = sess.run(predicts, feed_dict={image: np_img})

xmin, ymin, xmax, ymax, class_num = process_predicts(np_predict)

class_name = classes_name[class_num]

cv2.rectangle(resized_img, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0, 0, 255))

cv2.putText(resized_img, class_name, (int(xmin), int(ymin)), 2, 1.5, (0, 0, 255))

cv2.imwrite('cat_out.jpg', resized_img)

sess.close()