302 lines
12 KiB
Python
302 lines
12 KiB
Python
import datetime
|
|
import logging
|
|
import paho.mqtt.client as mqtt
|
|
|
|
import events.events_pb2
|
|
|
|
import depthai as dai
|
|
import cv2
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def to_tensor_result(packet):
|
|
return {
|
|
name: np.array(packet.getLayerFp16(name))
|
|
for name in [tensor.name for tensor in packet.getRaw().tensors]
|
|
}
|
|
|
|
|
|
def to_planar(frame):
|
|
return frame.transpose(2, 0, 1).flatten()
|
|
|
|
|
|
class FramePublisher:
|
|
def __init__(self, mqtt_client: mqtt.Client, frame_topic: str, img_width: int, img_height: int):
|
|
self._mqtt_client = mqtt_client
|
|
self._frame_topic = frame_topic
|
|
self._img_width = img_width
|
|
self._img_height = img_height
|
|
self._pipeline = self._configure_pipeline()
|
|
|
|
def _configure_pipeline(self) -> dai.Pipeline:
|
|
logger.info("configure pipeline")
|
|
pipeline = dai.Pipeline()
|
|
|
|
version = "2021.2"
|
|
pipeline.setOpenVINOVersion(version=dai.OpenVINO.Version.VERSION_2021_2)
|
|
|
|
# colorCam = pipeline.create(dai.node.ColorCamera)
|
|
# colorCam.setPreviewSize(256, 256)
|
|
# colorCam.setVideoSize(1024, 1024) # 4 times larger in both axis
|
|
# colorCam.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
|
|
# colorCam.setInterleaved(False)
|
|
# colorCam.setBoardSocket(dai.CameraBoardSocket.RGB)
|
|
# colorCam.setFps(10)
|
|
#
|
|
# controlIn = pipeline.create(dai.node.XLinkIn)
|
|
# controlIn.setStreamName('control')
|
|
# controlIn.out.link(colorCam.inputControl)
|
|
#
|
|
# cam_xout = pipeline.create(dai.node.XLinkOut)
|
|
# cam_xout.setStreamName('video')
|
|
# colorCam.video.link(cam_xout.input)
|
|
|
|
# ---------------------------------------
|
|
# 1st stage NN - text-detection
|
|
# ---------------------------------------
|
|
|
|
nn = pipeline.create(dai.node.NeuralNetwork)
|
|
nn.setBlobPath(
|
|
blobconverter.from_zoo(name="east_text_detection_256x256", zoo_type="depthai", shaves=6, version=version))
|
|
colorCam.preview.link(nn.input)
|
|
|
|
nn_xout = pipeline.create(dai.node.XLinkOut)
|
|
nn_xout.setStreamName('detections')
|
|
nn.out.link(nn_xout.input)
|
|
|
|
# ---------------------------------------
|
|
# 2nd stage NN - text-recognition-0012
|
|
# ---------------------------------------
|
|
|
|
manip = pipeline.create(dai.node.ImageManip)
|
|
manip.setWaitForConfigInput(True)
|
|
|
|
manip_img = pipeline.create(dai.node.XLinkIn)
|
|
manip_img.setStreamName('manip_img')
|
|
manip_img.out.link(manip.inputImage)
|
|
|
|
manip_cfg = pipeline.create(dai.node.XLinkIn)
|
|
manip_cfg.setStreamName('manip_cfg')
|
|
manip_cfg.out.link(manip.inputConfig)
|
|
|
|
manip_xout = pipeline.create(dai.node.XLinkOut)
|
|
manip_xout.setStreamName('manip_out')
|
|
|
|
nn2 = pipeline.create(dai.node.NeuralNetwork)
|
|
nn2.setBlobPath(blobconverter.from_zoo(name="text-recognition-0012", shaves=6, version=version))
|
|
nn2.setNumInferenceThreads(2)
|
|
manip.out.link(nn2.input)
|
|
manip.out.link(manip_xout.input)
|
|
|
|
nn2_xout = pipeline.create(dai.node.XLinkOut)
|
|
nn2_xout.setStreamName("recognitions")
|
|
nn2.out.link(nn2_xout.input)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
cam_rgb = pipeline.create(dai.node.ColorCamera)
|
|
xout_rgb = pipeline.create(dai.node.XLinkOut)
|
|
|
|
xout_rgb.setStreamName("rgb")
|
|
|
|
# Properties
|
|
cam_rgb.setBoardSocket(dai.CameraBoardSocket.RGB)
|
|
cam_rgb.setPreviewSize(width=self._img_width, height=self._img_height)
|
|
cam_rgb.setInterleaved(False)
|
|
cam_rgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.RGB)
|
|
cam_rgb.setFps(30)
|
|
|
|
# Linking
|
|
cam_rgb.preview.link(xout_rgb.input)
|
|
logger.info("pipeline configured")
|
|
return pipeline
|
|
|
|
def run(self):
|
|
|
|
with dai.Device(self._pipeline) as device:
|
|
q_vid = device.getOutputQueue("video", 4, blocking=False)
|
|
# This should be set to block, but would get to some extreme queuing/latency!
|
|
q_det = device.getOutputQueue("detections", 4, blocking=False)
|
|
|
|
q_rec = device.getOutputQueue("recognitions", 4, blocking=True)
|
|
|
|
q_manip_img = device.getInputQueue("manip_img")
|
|
q_manip_cfg = device.getInputQueue("manip_cfg")
|
|
q_manip_out = device.getOutputQueue("manip_out", 4, blocking=False)
|
|
|
|
controlQueue = device.getInputQueue('control')
|
|
|
|
frame = None
|
|
cropped_stacked = None
|
|
rotated_rectangles = []
|
|
rec_pushed = 0
|
|
rec_received = 0
|
|
host_sync = HostSeqSync()
|
|
|
|
characters = '0123456789abcdefghijklmnopqrstuvwxyz#'
|
|
codec = CTCCodec(characters)
|
|
|
|
ctrl = dai.CameraControl()
|
|
ctrl.setAutoFocusMode(dai.CameraControl.AutoFocusMode.CONTINUOUS_VIDEO)
|
|
ctrl.setAutoFocusTrigger()
|
|
controlQueue.send(ctrl)
|
|
|
|
while True:
|
|
vid_in = q_vid.tryGet()
|
|
if vid_in is not None:
|
|
host_sync.add_msg(vid_in)
|
|
|
|
# Multiple recognition results may be available, read until queue is empty
|
|
while True:
|
|
in_rec = q_rec.tryGet()
|
|
if in_rec is None:
|
|
break
|
|
rec_data = bboxes = np.array(in_rec.getFirstLayerFp16()).reshape(30, 1, 37)
|
|
decoded_text = codec.decode(rec_data)[0]
|
|
pos = rotated_rectangles[rec_received]
|
|
print("{:2}: {:20}".format(rec_received, decoded_text),
|
|
"center({:3},{:3}) size({:3},{:3}) angle{:5.1f} deg".format(
|
|
int(pos[0][0]), int(pos[0][1]), pos[1][0], pos[1][1], pos[2]))
|
|
# Draw the text on the right side of 'cropped_stacked' - placeholder
|
|
if cropped_stacked is not None:
|
|
cv2.putText(cropped_stacked, decoded_text,
|
|
(120 + 10, 32 * rec_received + 24),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
|
|
cv2.imshow('cropped_stacked', cropped_stacked)
|
|
rec_received += 1
|
|
|
|
if cv2.waitKey(1) == ord('q'):
|
|
break
|
|
|
|
if rec_received >= rec_pushed:
|
|
in_det = q_det.tryGet()
|
|
if in_det is not None:
|
|
frame = host_sync.get_msg(in_det.getSequenceNum()).getCvFrame().copy()
|
|
|
|
scores, geom1, geom2 = to_tensor_result(in_det).values()
|
|
scores = np.reshape(scores, (1, 1, 64, 64))
|
|
geom1 = np.reshape(geom1, (1, 4, 64, 64))
|
|
geom2 = np.reshape(geom2, (1, 1, 64, 64))
|
|
|
|
bboxes, confs, angles = east.decode_predictions(scores, geom1, geom2)
|
|
boxes, angles = east.non_max_suppression(np.array(bboxes), probs=confs, angles=np.array(angles))
|
|
rotated_rectangles = [
|
|
east.get_cv_rotated_rect(bbox, angle * -1)
|
|
for (bbox, angle) in zip(boxes, angles)
|
|
]
|
|
|
|
rec_received = 0
|
|
rec_pushed = len(rotated_rectangles)
|
|
if rec_pushed:
|
|
print("====== Pushing for recognition, count:", rec_pushed)
|
|
cropped_stacked = None
|
|
for idx, rotated_rect in enumerate(rotated_rectangles):
|
|
# Detections are done on 256x256 frames, we are sending back 1024x1024
|
|
# That's why we multiply center and size values by 4
|
|
rotated_rect[0][0] = rotated_rect[0][0] * 4
|
|
rotated_rect[0][1] = rotated_rect[0][1] * 4
|
|
rotated_rect[1][0] = rotated_rect[1][0] * 4
|
|
rotated_rect[1][1] = rotated_rect[1][1] * 4
|
|
|
|
# Draw detection crop area on input frame
|
|
points = np.int0(cv2.boxPoints(rotated_rect))
|
|
print(rotated_rect)
|
|
cv2.polylines(frame, [points], isClosed=True, color=(255, 0, 0), thickness=1,
|
|
lineType=cv2.LINE_8)
|
|
|
|
# TODO make it work taking args like in OpenCV:
|
|
# rr = ((256, 256), (128, 64), 30)
|
|
rr = dai.RotatedRect()
|
|
rr.center.x = rotated_rect[0][0]
|
|
rr.center.y = rotated_rect[0][1]
|
|
rr.size.width = rotated_rect[1][0]
|
|
rr.size.height = rotated_rect[1][1]
|
|
rr.angle = rotated_rect[2]
|
|
cfg = dai.ImageManipConfig()
|
|
cfg.setCropRotatedRect(rr, False)
|
|
cfg.setResize(120, 32)
|
|
# Send frame and config to device
|
|
if idx == 0:
|
|
w, h, c = frame.shape
|
|
imgFrame = dai.ImgFrame()
|
|
imgFrame.setData(to_planar(frame))
|
|
imgFrame.setType(dai.ImgFrame.Type.BGR888p)
|
|
imgFrame.setWidth(w)
|
|
imgFrame.setHeight(h)
|
|
q_manip_img.send(imgFrame)
|
|
else:
|
|
cfg.setReusePreviousImage(True)
|
|
q_manip_cfg.send(cfg)
|
|
|
|
# Get manipulated image from the device
|
|
transformed = q_manip_out.get().getCvFrame()
|
|
|
|
rec_placeholder_img = np.zeros((32, 200, 3), np.uint8)
|
|
transformed = np.hstack((transformed, rec_placeholder_img))
|
|
if cropped_stacked is None:
|
|
cropped_stacked = transformed
|
|
else:
|
|
cropped_stacked = np.vstack((cropped_stacked, transformed))
|
|
|
|
if cropped_stacked is not None:
|
|
cv2.imshow('cropped_stacked', cropped_stacked)
|
|
|
|
if frame is not None:
|
|
cv2.imshow('frame', frame)
|
|
|
|
key = cv2.waitKey(1)
|
|
if key == ord('q'):
|
|
break
|
|
elif key == ord('t'):
|
|
print("Autofocus trigger (and disable continuous)")
|
|
ctrl = dai.CameraControl()
|
|
ctrl.setAutoFocusMode(dai.CameraControl.AutoFocusMode.AUTO)
|
|
ctrl.setAutoFocusTrigger()
|
|
controlQueue.send(ctrl)
|
|
|
|
|
|
|
|
# Connect to device and start pipeline
|
|
with dai.Device(self._pipeline) as device:
|
|
logger.info('MxId: %s', device.getDeviceInfo().getMxId())
|
|
logger.info('USB speed: %s', device.getUsbSpeed())
|
|
logger.info('Connected cameras: %s', device.getConnectedCameras())
|
|
|
|
logger.info("output queues found: %s", device.getOutputQueueNames())
|
|
|
|
device.startPipeline()
|
|
# Queues
|
|
queue_size = 4
|
|
q_rgb = device.getOutputQueue("rgb", maxSize=queue_size, blocking=False)
|
|
|
|
while True:
|
|
try:
|
|
logger.debug("wait for new frame")
|
|
inRgb = q_rgb.get() # blocking call, will wait until a new data has arrived
|
|
|
|
im_resize = inRgb.getCvFrame()
|
|
|
|
is_success, im_buf_arr = cv2.imencode(".jpg", im_resize)
|
|
byte_im = im_buf_arr.tobytes()
|
|
|
|
now = datetime.datetime.now()
|
|
frame_msg = events.events_pb2.FrameMessage()
|
|
frame_msg.id.name = "robocar-oak-camera-oak"
|
|
frame_msg.id.id = str(int(now.timestamp() * 1000))
|
|
frame_msg.id.created_at.FromDatetime(now)
|
|
frame_msg.frame = byte_im
|
|
|
|
logger.debug("publish frame event to %s", self._frame_topic)
|
|
self._mqtt_client.publish(topic=self._frame_topic,
|
|
payload=frame_msg.SerializeToString(),
|
|
qos=0,
|
|
retain=False)
|
|
|
|
except Exception as e:
|
|
logger.exception("unexpected error: %s", str(e))
|