from handler.handler import Handler
from utils import visualization_utils as vis_util
import numpy as np
import cv2
import PIL.Image as Image


class DisplayHandler(Handler):

    def __init__(self, category_index):
        self.category_index = category_index

    def handle(self, frame, boxes, scores, classes, num):
        img = frame.image.copy()
        boxes = np.squeeze(boxes)
        scores = np.squeeze(scores)
        classes = np.squeeze(classes).astype(np.int32)
        for i in range(len(boxes)):
            if scores[i] < 0.6:
                continue;
            ymin, xmin, ymax, xmax = boxes[i]
            im_width, im_height = Image.fromarray(np.uint8(img)).convert('RGB').size
            (left, right, top, bottom) = (int(xmin * im_width), int(xmax * im_width),
                                          int(ymin * im_height), int(ymax * im_height))
            cv2.rectangle(img, (left, top), (right, bottom), (36,255,12), 2)
            #cv2.putText(img, self.category_index[classes[i]]["name"], (left - 10, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2, cv2.LINE_AA)

        #vis_util.visualize_boxes_and_labels_on_image_array(
        #   img,
        #   boxes,
        #   classes,
        #   scores,
        #   self.category_index,
        #   use_normalized_coordinates=True,
        #   line_thickness=8,
        #   min_score_thresh=0.60)

        cv2.imshow('Object detector', img)
        super().handle(frame, boxes, scores, classes, num)