initial commit

yolo_server/yolo_infer_test.py

@@ -0,0 +1,84 @@
+# 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)