class Listen2Stream(Routine):
def __init__(self, stream_address, queue, fps=30., *args, **kwargs):
super().__init__(*args, **kwargs)
self.stream_address = stream_address
self.is_file = str(stream_address).endswith("mp4")
self.stream = None
# self.stream = cv2.VideoCapture(self.stream_address)
self.q_handler = QueueHandler(queue)
self.fps = fps
self.updated_config = {}
def begin_capture(self):
self.stream = cv2.VideoCapture(self.stream_address)
if self.is_file:
self.fps = self.stream.get(cv2.CAP_PROP_FPS)
self.stream.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
self.stream.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
_, _ = self.grab_frame()
self.logger.info("Starting video capture on %s", self.stream_address)
def change_stream(self):
if self.stream_address == self.updated_config['stream_address']:
return
self.stream_address = self.updated_config['stream_address']
self.fps = self.updated_config['FPS']
self.is_file = str(self.stream_address).endswith("mp4")
self.logger.info("Changing source stream address to %s",
self.updated_config['stream_address'])
self.begin_capture()
def grab_frame(self):
grabbed, frame = self.stream.read()
msg = Message(frame, self.stream_address)
msg.record_entry(self.component_name, self.logger)
return grabbed, msg
def main_logic(self, *args, **kwargs):
if self.updated_config:
self.change_stream()
self.updated_config = {}
grabbed, msg = self.grab_frame()
if grabbed:
frame = msg.get_payload()
frame = resize(frame, 640, 480)
# if the stream is from a webcam, flip the frame
if self.stream_address == 0:
frame = cv2.flip(frame, 1)
msg.update_payload(frame)
success = self.q_handler.deque_non_blocking_put(msg)
return success
def setup(self, *args, **kwargs):
self.begin_capture()
def cleanup(self, *args, **kwargs):
self.stream.release()
del self.stream
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
class MessageFromRedis(Routine):
def __init__(self, in_key, url, queue, most_recent=True, *args, **kwargs):
super().__init__(*args, **kwargs)
self.in_key = in_key
self.url = url
self.q_handler = QueueHandler(queue)
self.msg_handler = None
self.most_recent = most_recent
self.read_method = None
self.flip = False
self.negative = False
def main_logic(self, *args, **kwargs):
encoded_msg = self.read_method(self.in_key)
if encoded_msg:
msg = message_decode(encoded_msg)
msg.record_entry(self.component_name, self.logger)
success = self.q_handler.deque_non_blocking_put(msg)
return success
else:
time.sleep(0)
return None
def setup(self, *args, **kwargs):
self.msg_handler = RedisHandler(self.url)
if self.most_recent:
self.read_method = self.msg_handler.read_most_recent_msg
else:
self.read_method = self.msg_handler.read_next_msg
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"))
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 MetaAndFrameFromRedis(Routine):
def __init__(self, in_key_meta, in_key_im, url, queue, *args, **kwargs):
super().__init__(*args, **kwargs)
self.in_key_meta = in_key_meta
self.in_key_im = in_key_im
self.url = url
self.q_handler = QueueHandler(queue)
self.msg_handler = None
self.flip = False
self.negative = False
def receive_msg(self, in_key, most_recent=True):
if most_recent:
encoded_msg = self.msg_handler.read_most_recent_msg(in_key)
else:
encoded_msg = self.msg_handler.receive(in_key)
if not encoded_msg:
return None
msg = message_decode(encoded_msg)
msg.record_entry(self.component_name, self.logger)
return msg
def main_logic(self, *args, **kwargs):
pred_msg = self.receive_msg(self.in_key_meta, most_recent=False)
frame_msg = self.receive_msg(self.in_key_im, most_recent=True)
if frame_msg:
arr = frame_msg.get_payload()
if self.flip:
arr = cv2.flip(arr, 1)
if self.negative:
arr = 255 - arr
frame_msg.update_payload(arr)
success = self.q_handler.deque_non_blocking_put(
(frame_msg, pred_msg))
return success
else:
time.sleep(0)
return False
def setup(self, *args, **kwargs):
self.msg_handler = RedisHandler(self.url)
def cleanup(self, *args, **kwargs):
self.msg_handler.close()
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 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