class Message2Redis(Routine):
def __init__(self, out_key, url, queue, maxlen, *args, **kwargs):
super().__init__(*args, **kwargs)
self.out_key = out_key
self.url = url
self.q_handler = QueueHandler(queue)
self.maxlen = maxlen
self.msg_handler = None
def main_logic(self, *args, **kwargs):
msg = self.q_handler.non_blocking_get()
if msg:
msg.record_exit(self.component_name, self.logger)
encoded_msg = message_encode(msg)
self.msg_handler.send(self.out_key, encoded_msg)
return True
else:
return False
def setup(self, *args, **kwargs):
self.msg_handler = RedisHandler(self.url, self.maxlen)
self.msg_handler.connect()
def cleanup(self, *args, **kwargs):
self.msg_handler.close()
class VisLogic(Routine):
def __init__(self, in_queue, out_queue, *args, **kwargs):
super().__init__(*args, **kwargs)
self.in_queue = QueueHandler(in_queue)
self.out_queue = QueueHandler(out_queue)
self.vis = VideoVisualizer(MetadataCatalog.get("coco_2017_train"))
def main_logic(self, *args, **kwargs):
# TODO implement input that takes both frame and metadata
messages = self.in_queue.non_blocking_get()
if messages:
frame_msg, pred_msg = messages
if pred_msg is not None and not pred_msg.is_empty():
frame = frame_msg.get_payload()
pred = pred_msg.get_payload()
image = self.vis.draw_instance_predictions(frame, pred) \
.get_image()
frame_msg.update_payload(image)
frame_msg.history = pred_msg.history
frame_msg.record_exit(self.component_name, self.logger)
success = self.out_queue.deque_non_blocking_put(frame_msg)
return success
else:
return None
def setup(self, *args, **kwargs):
self.state.dropped = 0
def cleanup(self, *args, **kwargs):
pass
def gen(q: QueueHandler):
while True:
encoded_frame = q.non_blocking_get()
if encoded_frame:
yield (b'--frame\r\n'
b'Pragma-directive: no-cache\r\n'
b'Cache-directive: no-cache\r\n'
b'Cache-control: no-cache\r\n'
b'Pragma: no-cache\r\n'
b'Expires: 0\r\n'
b'Content-Type: image/jpeg\r\n\r\n' + encoded_frame +
b'\r\n\r\n')
def gen(q: QueueHandler):
while True:
msg = q.non_blocking_get()
if msg:
image = msg.get_payload()
ret, frame = cv2.imencode('.jpg', image)
frame = frame.tobytes()
yield (b'--frame\r\n'
b'Pragma-directive: no-cache\r\n'
b'Cache-directive: no-cache\r\n'
b'Cache-control: no-cache\r\n'
b'Pragma: no-cache\r\n'
b'Expires: 0\r\n'
b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')
class VisLogic(Routine):
def __init__(self, in_queue, out_queue, *args, **kwargs):
super().__init__(*args, **kwargs)
self.in_queue = QueueHandler(in_queue)
self.out_queue = QueueHandler(out_queue)
self.vis = VideoVisualizer(MetadataCatalog.get("coco_2017_train"))
self.latest_pred_msg = None
def main_logic(self, *args, **kwargs):
# TODO implement input that takes both frame and metadata
messages = self.in_queue.non_blocking_get()
if messages:
frame_msg, pred_msg = messages
if pred_msg is not None:
self.latest_pred_msg = pred_msg
pred_msg = self.latest_pred_msg
self.draw_preds_on_frame(frame_msg, pred_msg)
self.pass_frame_to_flask(frame_msg, pred_msg)
return True
else:
return None
def draw_preds_on_frame(self, frame_msg, pred_msg: Optional[Message]):
if pred_msg is not None and not pred_msg.is_empty():
frame = frame_msg.get_payload()
pred = pred_msg.get_payload()
image = self.vis.draw_instance_predictions(frame, pred) \
.get_image()
frame_msg.update_payload(image)
def pass_frame_to_flask(self, frame_msg, pred_msg: Optional[Message]):
image = frame_msg.get_payload()
_, frame = cv2.imencode('.jpg', image)
frame = frame.tobytes()
frame_msg.record_exit(self.component_name, self.logger)
if pred_msg is not None and not pred_msg.reached_exit:
pred_msg.record_exit(self.component_name, self.logger)
latency = pred_msg.get_end_to_end_latency(self.component_name)
self.metrics_collector.collect_latency(latency,
self.component_name)
success = self.out_queue.deque_non_blocking_put(frame)
return success
def setup(self, *args, **kwargs):
self.state.dropped = 0
def cleanup(self, *args, **kwargs):
pass
class VideoWriterLogic(Routine):
def __init__(self,
output_file,
in_queue,
fps=30,
im_size=(640, 480),
*args,
**kwargs):
super().__init__(*args, **kwargs)
self.writer = None
self.output_file = output_file
self.q_handler = QueueHandler(in_queue)
self.w, self.h = im_size
self.fps = fps
def main_logic(self, *args, **kwargs):
msg: Message = self.q_handler.non_blocking_get()
if msg:
frame = msg.get_payload()
if frame.shape[0] != self.h or frame.shape[1] != self.w:
self.writer.release()
self.writer = cv2.VideoWriter(
self.output_file,
cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), self.fps,
(frame.shape[1], frame.shape[0]))
self.h, self.w = frame.shape[0], frame.shape[1]
frame = cv2.putText(frame, f"{msg.id.split('_')[-1]}",
(0, self.h - 10), cv2.FONT_HERSHEY_SIMPLEX, 1,
(0, 0, 255), 2, cv2.LINE_AA)
self.writer.write(frame)
def setup(self, *args, **kwargs):
self.writer = cv2.VideoWriter(
self.output_file, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'),
self.fps, (self.w, self.h))
def cleanup(self, *args, **kwargs):
self.writer.release()
class SORTLogic(Routine):
def __init__(self, in_queue, out_queue, component_name, *args, **kwargs):
super().__init__(component_name=component_name)
self.in_queue = QueueHandler(in_queue)
self.out_queue = QueueHandler(out_queue)
self.sort = InstancesSort(*args, **kwargs)
def main_logic(self, *args, **kwargs):
pred_msg = self.in_queue.non_blocking_get()
if pred_msg:
instances = pred_msg.get_payload()
new_instances = self.sort.update_instances(instances)
pred_msg.update_payload(new_instances)
success = self.out_queue.deque_non_blocking_put(pred_msg)
return success
else:
return False
def setup(self, *args, **kwargs):
self.state.dropped = 0
def cleanup(self, *args, **kwargs):
pass
class PredsToDatabase(Routine):
def __init__(self, in_queue, db_conf, *args, **kwargs):
super().__init__(*args, **kwargs)
self.queue = QueueHandler(in_queue)
self.logger.addHandler(logging.StreamHandler())
self.add_event_handler(Events.AFTER_LOGIC, self.commit_to_db)
self.dbh = MarsDBHandler(db_conf, logger=self.logger)
self.is_on = None
def setup(self, *args, **kwargs):
# make sure connection to db is up and turn on
self.toggle()
def main_logic(self, *args, **kwargs):
msg = self.queue.non_blocking_get()
# Only do something if msg isn't empty and saving is toggled on
if self.is_on and msg:
msg_id = msg.id.split("_")[-1]
timestamp = dt.fromtimestamp(msg.history["VideoCapture"]["entry"])
for prediction in msg.get_payload():
# get the needed fields and convert to a format that is good to be inserted into the db
box = prediction.pred_boxes.tensor.numpy().squeeze().astype(
int) if prediction.has("pred_boxes") else None
objectness = prediction.scores.numpy().item(
) if prediction.has("scores") else None
class_scores = prediction.class_scores.numpy().squeeze(
) if prediction.has("class_scores") else None
# convert the list of class scores into a dict. astype(float) since np.float32 is not JSON-Serializable
score_dict = {
idx: val.astype(float)
for idx, val in enumerate(class_scores)
}
# create the object and add to session
pred = self.dbh.tables['predictions'](msg_id, box, objectness,
score_dict,
msg.source_address,
timestamp)
self.dbh.session.add(pred)
def cleanup(self, *args, **kwargs):
# make sure session is closed
self.dbh.session.close()
def toggle(self):
"""
Toggle saving to the database on or off. Makes sure db connection is up when
toggled ON.
Returns:
A message describing the new state (On/Off)
"""
if self.is_on:
self.is_on = False
else:
self.is_on = self.dbh.test_connection()
msg = "Saving predictions toggled " + ("ON" if self.is_on else "OFF")
self.logger.info(msg)
return msg
def commit_to_db(self, routine):
"""
Commit the current session to database, rollback in case of an error
Args:
routine: This is just a dummy argument, but its necessary if this method
is to be used as an event handler.
"""
if self.is_on:
try:
self.dbh.session.commit()
# if there was any Database related error
except SQLAlchemyError as sqlae:
self.dbh.session.rollback()
self.logger.error(
format_sqla_error(sqlae) + f"\nError code: {sqlae.code}")
# if there was an unknown error
except Exception:
self.dbh.session.close()
raise
class YoloV3Logic(Routine):
def __init__(self, in_queue, out_queue, *args, **kwargs):
super().__init__(*args, **kwargs)
self.in_queue = QueueHandler(in_queue)
self.out_queue = QueueHandler(out_queue)
self.img_size = (
320, 192
) if ONNX_EXPORT else opt.img_size # (320, 192) or (416, 256) or (608, 352)
out, source, weights, half = opt.output, opt.source, opt.weights, opt.half
device = torch_utils.select_device(force_cpu=ONNX_EXPORT)
self.model = Darknet(opt.cfg, self.img_size)
if weights.endswith('.pt'): # pytorch format
self.model.load_state_dict(
torch.load(weights, map_location=device)['model'])
else: # darknet format
_ = load_darknet_weights(self.model, weights)
self.model.fuse()
self.model.to(device).eval()
# Half precision
self.half = half and device.type != 'cpu' # half precision only supported on CUDA
if half:
self.model.half()
self.classes = load_classes(opt.names)
self.colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(self.classes))]
self.device = device
def main_logic(self, *args, **kwargs):
msg = self.in_queue.non_blocking_get()
if msg:
self.logger.info("Received the following message: %s", str(msg))
im0 = msg.get_payload()
img, *_ = letterbox(im0, new_shape=self.img_size)
# Normalize RGB
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB
img = np.ascontiguousarray(img,
dtype=np.float16 if self.half else
np.float32) # uint8 to fp16/fp32
img /= 255.0
img = torch.from_numpy(img).to(self.device)
# print(f"yolo {self.device}")
if img.ndimension() == 3:
img = img.unsqueeze(0)
with torch.no_grad():
pred, _ = self.model(img)
det = non_max_suppression(pred, opt.conf_thres, opt.nms_thres)[0]
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# print(det.shape)
# print(det)
# for *xyxy, conf, _, cls in det:
# label = '%s %.2f' % (self.classes[int(cls)], conf)
# plot_one_box(xyxy, im0, label=label, color=self.colors[int(cls)])
res = Instances(im0.shape)
res.set("pred_boxes", Boxes(det[:, :4]))
res.set("scores", det[:, 4])
res.set("class_scores", det[:, 5])
res.set("pred_classes", det[:, 6].round().int())
else:
res = Instances(im0.shape)
res.set("pred_boxes", [])
msg.payload = PredictionPayload(res.to("cpu"))
success = self.out_queue.deque_non_blocking_put(msg)
return success
else:
return None
def setup(self, *args, **kwargs):
self.state.dropped = 0
def cleanup(self, *args, **kwargs):
del self.model, self.device, self.classes, self.colors