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