2024-04-26 22:11:53 -07:00
|
|
|
|
|
|
|
|
import cv2
|
|
|
|
|
from websocket import create_connection
|
|
|
|
|
import win32file
|
2024-04-27 22:37:06 -07:00
|
|
|
import struct
|
2024-04-26 22:11:53 -07:00
|
|
|
import numpy as np
|
|
|
|
|
from ultralytics import YOLO
|
|
|
|
|
|
|
|
|
|
# object classes
|
|
|
|
|
classNames = ["person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat",
|
|
|
|
|
"traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
|
|
|
|
|
"dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
|
|
|
|
|
"handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat",
|
|
|
|
|
"baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
|
|
|
|
|
"fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
|
|
|
|
|
"carrot", "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", "bed",
|
|
|
|
|
"diningtable", "toilet", "tvmonitor", "laptop", "mouse", "remote", "keyboard", "cell phone",
|
|
|
|
|
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
|
|
|
|
|
"teddy bear", "hair drier", "toothbrush"
|
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def handle_connection(pipe):
|
|
|
|
|
# Load YOLO model
|
|
|
|
|
model = YOLO('yolov8s.pt') # Load an official Detect model
|
|
|
|
|
|
2024-04-27 22:37:06 -07:00
|
|
|
print("Pipe connected!")
|
|
|
|
|
try:
|
|
|
|
|
while True:
|
|
|
|
|
data = win32file.ReadFile(pipe, 640*480*3)[1]
|
|
|
|
|
|
|
|
|
|
if not data:
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
nparr = np.frombuffer(data, np.uint8).reshape((480, 640, 3))
|
|
|
|
|
|
|
|
|
|
# frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
|
|
|
|
frame = cv2.cvtColor(nparr, cv2.COLOR_BGR2RGB)
|
|
|
|
|
|
|
|
|
|
# Perform object detection
|
|
|
|
|
results = model.track(frame, persist=True)
|
|
|
|
|
|
|
|
|
|
for r in results:
|
|
|
|
|
lines = ""
|
|
|
|
|
boxes = r.boxes
|
|
|
|
|
for box in boxes:
|
|
|
|
|
if box.cls[0].item() == 0 and not box.id is None:
|
|
|
|
|
# bounding box
|
|
|
|
|
id = box.id.int().cpu().tolist()
|
|
|
|
|
# x1, y1, x2, y2 = box.xyxy[0]
|
|
|
|
|
# # 2.08333 = 1/480 * 1000 or normalize, then save 4 sig-figures and cast to int
|
|
|
|
|
# # 1.5625 = 1/640 * 1000 or normalize, then save 4 sig-figures and cast to int
|
|
|
|
|
|
|
|
|
|
# x1, x2, y1, y2 = int(x1), int(x2), int(y1), int(y2)
|
|
|
|
|
|
|
|
|
|
# # put box in cam
|
|
|
|
|
# cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 255), 3)
|
|
|
|
|
|
|
|
|
|
# # class name
|
|
|
|
|
# cls = int(box.cls[0])
|
|
|
|
|
|
|
|
|
|
# # object details
|
|
|
|
|
# org = [x1, y1]
|
|
|
|
|
# org2 = [x1, y1+50]
|
|
|
|
|
# font = cv2.FONT_HERSHEY_SIMPLEX
|
|
|
|
|
# fontScale = 1
|
|
|
|
|
# color = (255, 0, 0)
|
|
|
|
|
# color_w = (255, 255, 255)
|
|
|
|
|
# thickness = 2
|
|
|
|
|
|
|
|
|
|
# cv2.putText(frame, classNames[cls], org, font, fontScale, color, thickness)
|
|
|
|
|
# cv2.putText(frame, str(id), org2, font, fontScale, color_w, thickness)
|
|
|
|
|
|
|
|
|
|
x1, y1, x2, y2 = box.xyxyn[0]
|
|
|
|
|
x1, y1, x2, y2 = int(x1 * 1000), int(y1 * 1000), int(x2 * 1000), int(y2 * 1000)
|
|
|
|
|
lines += f"{id} {x1}:{y1} {x2}:{y2}\n"
|
|
|
|
|
|
|
|
|
|
# cv2.imshow('Webcam', frame)
|
|
|
|
|
# if cv2.waitKey(1) == ord('q'):
|
|
|
|
|
# break
|
|
|
|
|
|
|
|
|
|
ret = lines.encode('utf-8')
|
|
|
|
|
length_prefix = struct.pack('I', len(ret))
|
|
|
|
|
|
|
|
|
|
win32file.WriteFile(pipe, length_prefix+ret)
|
|
|
|
|
except Exception as e:
|
|
|
|
|
print("Pipe closed with error: " + str(e))
|
|
|
|
|
return
|
|
|
|
|
# finally:
|
|
|
|
|
# cv2.destroyAllWindows()
|
