85 lines
2.4 KiB
Python
85 lines
2.4 KiB
Python
# 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)
|