# from infer import Yolo_model_infer
# import cv2


# infer = Yolo_model_infer()

# image = cv2.imread("ball_0094.png")
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# results = infer.infer(image)

# expect_boxes = (results[:, 1] > 0.5) & (results[:, 0] > -1)
# np_boxes = results[expect_boxes, :]

# print(np_boxes)

from infer_new import Yolo_model_infer
import cv2
from visualize import visualize_box_mask
import zmq
import numpy as np
from loguru import logger

import time


# infer = Yolo_model_infer()
# image = cv2.imread("20240108161722.jpg")
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# results = infer.infer(image)
# print(results)
# expect_boxes = (results[:, 1] > 0.5) & (results[:, 0] > -1)
# np_boxes = results[expect_boxes, :]
# print(np_boxes)

# context = zmq.Context()
# camera1_socket = context.socket(zmq.SUB)
# hwm = 5
# camera1_socket.setsockopt(zmq.RCVHWM, hwm) 
# camera1_socket.connect("tcp://localhost:5556")
# camera1_socket.setsockopt_string(zmq.SUBSCRIBE, "")
 
# camera1_socket.set_hwm(1)
context1 = zmq.Context()
socket_server = context1.socket(zmq.PUB)
socket_server.bind("tcp://*:7777")


labels = [
    "tower", "sign", "shelter", "hospital", "basket", "base", 
    "Yball", "Spiller", "Rmark", "Rblock", "Rball", "Mpiller", 
    "Lpiller", "Lmark", "Bblock", "Bball"
]

infer = Yolo_model_infer()
cap = cv2.VideoCapture(2)
cap.set(cv2.CAP_PROP_FRAME_WIDTH,320)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT,240)
ret = True
while True:
    # message = camera1_socket.recv()
    # np_array = np.frombuffer(message, dtype=np.uint8)
    # frame = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
    ret, frame = cap.read()
    
    if ret:

        results = infer.infer(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        # logger.info("111")
        img = visualize_box_mask(frame,results,labels)
        showim = np.array(img)
        # cv2.imshow("Received", showim)
        _, encode_frame = cv2.imencode(".jpg", showim)
        socket_server.send(encode_frame.tobytes())
    # if cv2.waitKey(1) == 27:
    #     break
    # image = cv2.imread("20240525_170248.jpg")
    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # results = infer.infer(image)

#     # expect_boxes = (results[:, 1] > 0.5) & (results[:, 0] > -1)
#     # np_boxes = results[expect_boxes, :]
#     # print(np_boxes)
#     # img = visualize_box_mask(image,results,labels)
#     # img.save('20240525_170248_box.jpg', quality=95)