import cv2
import time
import numpy as np
import core.utils as utils
import tensorflow as tf
from PIL import Image,ImageGrab

def cap_demo():
    # 模型pb文件路径
    pb_file = "./yolov3_coco.pb"
    # 目标检测类别总数
    num_classes = 80
    # 输入图像的尺寸
    input_size = 416

    graph = tf.Graph()
    return_elements = ["input/input_data:0", "pred_sbbox/concat_2:0", "pred_mbbox/concat_2:0", "pred_lbbox/concat_2:0"]
    return_tensors = utils.read_pb_return_tensors(graph, pb_file, return_elements)
    with tf.Session(graph=graph) as sess:
        while True:
            img=ImageGrab.grab()
            frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
            frame_size = frame.shape[:2]
            image_data = utils.image_preporcess(np.copy(frame), [input_size, input_size])
            image_data = image_data[np.newaxis, ...]
            pred_sbbox, pred_mbbox, pred_lbbox = sess.run(
                [return_tensors[1], return_tensors[2], return_tensors[3]],
                    feed_dict={return_tensors[0]: image_data})
            pred_bbox = np.concatenate([np.reshape(pred_sbbox, (-1, 5 + num_classes)),
                                    np.reshape(pred_mbbox, (-1, 5 + num_classes)),
                                    np.reshape(pred_lbbox, (-1, 5 + num_classes))], axis=0)

            bboxes = utils.postprocess_boxes(pred_bbox, frame_size, input_size, 0.3)
            # 非极大值抑制，IOU的阈值设为0.45
            bboxes = utils.nms(bboxes, 0.45, method='nms')
            # 得到的结果是一张张图片
            image = utils.draw_bbox(frame, bboxes)
            cv2.namedWindow("result", cv2.WINDOW_AUTOSIZE)
            result = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            cv2.imshow('result',result)
            print("执行了")
            if cv2.waitKey(1) & 0xFF == ord('q'): 
                cv2.destroyAllWindows()
                break

if __name__=="__main__":
    cap_demo()