def test_subscribe(self):
queue = Queue()
test_delay_sec = 15
producer = Producer(self.config)
consumer = KafkaDelayConsumer(self.config, delay_sec=test_delay_sec)
print("start!!")
def task(record):
print("%s subscribe! %s" % (record.topic(), record.value()))
queue.put((record.topic(), int(record.value()) + 1, time.time()))
consumer.subscribe(['test'], task)
time.sleep(3)
producer.produce('test', str(20))
produce_time = time.time()
print("%d produce!" % (produce_time))
time.sleep(9)
producer.produce('test', str(30))
produce_time2 = time.time()
print("%d produce!" % (time.time()))
(topic, value, consume_time) = queue.get()
print("get message %d", consume_time - produce_time, consume_time)
#self.assertTrue(test_delay_sec <= int(consume_time - produce_time) <= test_delay_sec+1)
(topic, value, consume_time) = queue.get()
print("get message2 %d", consume_time - produce_time2, consume_time)
consumer.unsubscribe()
print("end")
def produce(self, topic, value=None, *args, **kwargs):
"""
overridden method
:param topic:
:param value:
:param args:
:param kwargs:
:return:
"""
msg_header_dict = dict() if kwargs.get('headers') is None else dict(
kwargs.get('headers'))
parent_context = self.tracer.extract(Format.TEXT_MAP, msg_header_dict)
user_tags = {}
for span_tag_provider in self.span_tags_providers:
user_tags = merge_two_dicts(
user_tags, span_tag_provider(topic, kwargs.get('key'), value))
span = self.tracer.start_span(
self.span_name_provider(topic, kwargs.get('key'), value),
references=[follows_from(parent_context)],
tags=merge_two_dicts(producer_span_tags_provider(topic),
user_tags))
# Inject created span context into message header for sending to kafka queue
self.tracer.inject(span.context, Format.TEXT_MAP, msg_header_dict)
kwargs['headers'] = list(msg_header_dict.items())
kwargs['on_delivery'] = create_tracing_delivery_callback(
kwargs['on_delivery'], span)
Producer.produce(self, topic, value, *args, **kwargs)
def test_plain_text_consume_and_produce(
consumer_group,
producer: ConfluenceProducer,
source_topic: Tuple[str, int],
target_topic: Tuple[str, int],
tmpdir_factory,
):
source_topic_id, _ = source_topic
target_topic_id, _ = target_topic
working_dir = tmpdir_factory.mktemp("working_directory")
produced_messages = produce_test_messages(producer, source_topic)
file_consumer = FileConsumer(consumer_group, source_topic_id, working_dir,
False)
file_consumer.consume(10)
producer = FileProducer(working_dir)
producer.produce(target_topic_id)
# Check assertions:
assertion_check_directory = tmpdir_factory.mktemp(
"assertion_check_directory")
file_consumer = FileConsumer((consumer_group + "assertion_check"),
target_topic_id, assertion_check_directory,
False)
file_consumer.consume(10)
consumed_messages = get_consumed_messages(assertion_check_directory, False)
assert produced_messages == consumed_messages
class KafkaProducer(BaseProducer):
def __init__(self, topic_name, num_partitions=NUM_PARTITIONS, num_replicas=NUM_REPLICAS):
BaseProducer.__init__(self, topic_name, num_partitions, num_replicas)
self.producer = Producer({'bootstrap.servers': BROKER_URL})
def produce(self, encoded_data):
self.producer.produce(self.topic_name, encoded_data)
class ConfluentKafkaProducer(KafkaProducerInterface):
def __init__(self, config, logger):
super(ConfluentKafkaProducer, self).__init__()
self.logger = logger
conf = {
'bootstrap.servers': ','.join(config['kafka_bootstrap_servers']),
'broker.version.fallback': config['kafka_broker_version_fallback'],
'api.version.request': config['kafka_api_version_request'],
'queue.buffering.max.ms': config['kafka_producer_batch_linger_ms'],
'queue.buffering.max.kbytes':
config['kafka_producer_buffer_kbytes'],
'message.send.max.retries': 3,
'default.topic.config': {
'request.required.acks': 1
}
}
self.logger.info("Creating a Confluent Kafka Producer",
{"config": json.dumps(conf, indent=4)})
self.producer = Producer(dict(conf,
**{'error_cb': self.error_callback}),
logger=logger.logger)
# Service any logging
self.producer.poll(0.25)
def error_callback(self, error):
"""
:param error:
:type error: KafkaError
:param message:
:param datum:
:return:
"""
if error:
datum = {}
datum['success'] = False
datum['exception'] = error.name()
datum['description'] = error.str()
self.logger.error("Kafka error", datum if datum else {})
def send(self, topic, message, callback=None):
self.producer.produce(topic,
json.dumps(message).encode('utf-8'),
callback=callback)
# Service the delivery callback queue.
self.producer.poll(0)
def poll(self):
self.producer.poll(0)
def close(self):
self.producer.flush()
self.producer.poll(0)
def produce(queue, key, val, headers=None):
logging.info("Producing into %s: %s %s", queue, key, val)
producer = Producer({
'bootstrap.servers': KAFK,
"message.send.max.retries": 2
})
producer.poll(0)
producer.produce(queue, key=key, value=val, headers=headers)
producer.flush()
def check_connection(self):
def check_callback(error, event):
if error:
if error.code() == KafkaError._MSG_TIMED_OUT:
log.error(
"This Timout might indicate the broker is down or connection is misconfigured"
)
log.error(f"Error while producing initial msg: {error}")
raise KafkaSetupError()
config = ConfigFactory(kafka_client="producer").config
config["delivery.timeout.ms"] = "3000" # 3 seconds
prod = Producer(config)
prod.produce("pyrandall", "starting simulate", callback=check_callback)
prod.flush() # block until callback is called
def produce_test_messages(producer: ConfluenceProducer,
topic: Tuple[str, int]) -> Iterable[KafkaMessage]:
topic_name, num_partitions = topic
messages = []
for i in range(10):
partition = random.randrange(0, num_partitions)
random_value = "".join(random.choices(ascii_letters, k=5))
message = KafkaMessage(str(i), random_value, partition)
messages.append(message)
producer.produce(topic=topic_name,
key=message.key,
value=message.value,
partition=message.partition)
producer.flush()
return messages
def produce():
try:
p = Producer({
'bootstrap.servers':
'192.168.150.6:9092,192.168.150.7:9092,192.168.150.8:9092'
})
# data = requests.get('http://hq.sinajs.cn/list=sh601006,sh601008')
# print(data.text)
# for dt in data.text.split(";"):
# d = dt.split("=")[1]
# ret = p.produce('ddc_test_topic', d.encode('utf-8'), callback=delivery_report)
# # print(ret)
for i in range(10000):
data = '[{"content": "clogan header", "contentType": 1, "curTime": 1587432527233, "threadName": "clogan", "threadId": "1", "mainThread": true}, {"content": "test logan 0", "contentType": 1, "curTime": 1587432527212, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432578763, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432578955, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432579125, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432579277, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432579440, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432579793, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432579954, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432580091, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432580234, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432580369, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432580526, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432580663, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432580804, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432580927, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432581045, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432581201, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432581333, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432581453, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432581584, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432581720, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432581850, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432581984, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432582327, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432582485, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432582641, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432582792, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432582941, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432583098, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432583256, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432583408, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432583558, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432583714, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432583871, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432584018, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432584164, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432584300, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432584532, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432584691, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432584835, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432584991, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432585156, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432585321, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432585469, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432585605, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432585769, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432585926, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432586070, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432586226, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432586377, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432586533, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432586691, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432586842, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432586998, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432587150, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432587297, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432587450, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432587593, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432587735, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432587841, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432587970, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432588112, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432588227, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432588363, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432588498, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432588636, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432588770, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432588920, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432589048, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432589207, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432589372, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432589523, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432589671, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432589835, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432589986, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432590156, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432590322, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432590472, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432590622, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432590770, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432590913, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432591049, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432591179, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432591336, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432591492, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432591668, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432591810, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432591971, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432592136, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432592285, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432592442, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432592602, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432592749, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432592914, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432593067, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432593275, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432593429, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432593575, "threadName": "main", "threadId": "2", "mainThread": true}, {"content": "test test testtes ", "contentType": 11, "curTime": 1587432593708, "threadName": "main", "threadId": "2", "mainThread": true}]'
# Trigger any available delivery report callbacks from previous produce() calls
print(len(data))
p.poll(0)
#
# # Asynchronously produce a message, the delivery report callback
# # will be triggered from poll() above, or flush() below, when the message has
# # been successfully delivered or failed permanently.
ret = p.produce('Applog',
data.encode('utf-8'),
callback=delivery_report)
print(ret)
# Wait for any outstanding messages to be delivered and delivery report
# callbacks to be triggered.
# p.flush()
except Exception as e:
print(traceback.format_exc())
class EventsProducer:
def __init__(self):
self.loop = asyncio.get_event_loop()
self.loop.run_until_complete(self.__start_producer())
def send_upsert_chat_entity(self, members, lab_id, chat_id, name):
event = create_event(entity={"containerReferenceId": lab_id, "chatId": chat_id, "name": name, "members": members})
self.__send_event(UPSERT_CHAT_ENTITY, event)
def send_delete_chat_entity(self, chat_id):
event = create_event(entity={"chatId": chat_id})
self.__send_event(DELETE_CHAT_ENTITY, event)
def send_add_chat_member(self, user_auth_key, chat_id):
event = create_event(chat_id=chat_id, entity={"userAuthKey": user_auth_key, "chatId": chat_id})
self.__send_event(ADD_CHAT_MEMBER, event)
def send_remove_chat_member(self, user_auth_key, chat_id):
event = create_event(entity={"userAuthKey": user_auth_key, "chatId": chat_id})
self.__send_event(REMOVE_MEMBERS_FROM_CHAT_ENTITY, event)
async def __start_producer(self):
self.producer = Producer(({
"bootstrap.servers": settings.BOOTSTRAP_SERVERS,
"security.protocol": "SASL_SSL",
"sasl.mechanisms": "PLAIN",
"sasl.username": settings.SASL_PLAIN_USERNAME,
"sasl.password": settings.SASL_PLAIN_PASSWORD
}))
def __send_event(self, key, event: dict):
self.loop.run_until_complete(self.__send_one(key, event))
def __delivery_report(self, err, msg):
""" Called once for each message produced to indicate delivery result.
Triggered by poll() or flush(). """
if err is not None:
print('Message delivery failed: {}'.format(err))
else:
print('Message delivered to {} [{}]'.format(msg.topic(), msg.partition()))
async def __send_one(self, key, event: dict):
self.producer.produce(topic=settings.TOPIC_CHAT, value=json.dumps(event), key=key,
callback=self.__delivery_report)
self.producer.poll(10000)
class StopEventsDataPublisher:
_logger = logging.getLogger('StopEventsDataPublisher')
def __init__(self):
kafka_configs = KafkaHelper.get_kafka_configs()
self._producer = Producer(kafka_configs)
def publish_stop_event_records(self,
stop_event_records,
topic=STOP_EVENT_TOPIC):
self._logger.info(
"Publishing {} stop event records to {} topic ...".format(
len(stop_event_records), topic))
delivered_records = 0
def callback(err, msg):
nonlocal delivered_records
if err is not None:
self._logger.error("Failed to deliver message: %s: %s" %
(str(msg), str(err)))
else:
delivered_records += 1
self._logger.debug(
"Published record to topic {} partition [{}] @ offset {}".
format(msg.topic(), msg.partition(), msg.offset()))
self._logger.debug(
'Published records count: {}'.format(delivered_records))
for trip_id, stop_events in stop_event_records.items():
stop_data_record = dict()
stop_data_record[trip_id] = stop_events
self._producer.produce(topic,
value=json.dumps(stop_data_record),
on_delivery=callback)
self._producer.poll(0)
self._producer.flush()
self._logger.info(
'Done delivering records to {} topic! A total of {} records were published'
.format(topic, delivered_records))
def test_subscribe(self):
test_retry_backoff = [3, 13]
queue = Queue()
config = KafkaConfig(bootstrap_servers='127.0.0.1:9093',
group_id='testConsumer%s' % (random.randint),
auto_offset_reset='largest')
producer = Producer(config)
consumer = KafkaRetryConsumer(config, retry_backoff=test_retry_backoff)
print("start!!")
def task(record):
print("[%20s] subscribe! %s" % (record.topic(), record.value()))
queue.put((record.topic(), int(record.value()) + 1, time.time()))
raise KafkaRetryException
consumer.subscribe(['test'], task)
try:
print("here!!")
time.sleep(10)
producer.produce('test', str(30))
produce_time = time.time()
print("%d produce!" % (produce_time))
time.sleep(2)
producer.produce('test', str(50))
produce_time2 = time.time()
print("%d produce2!" % (produce_time))
for i in range(6):
(topic, value, consume_time) = queue.get()
if value is 31:
print("30 get message %d", consume_time - produce_time)
else:
print("50 get message %d", consume_time - produce_time2)
finally:
consumer.unsubscribe()
print("end")
def test_plain_text_consume_and_produce(
consumer_group,
producer: ConfluenceProducer,
source_topic: Tuple[str, int],
target_topic: Tuple[str, int],
tmpdir_factory,
):
source_topic_id, _ = source_topic
target_topic_id, _ = target_topic
output_directory = tmpdir_factory.mktemp("output_directory")
produced_messages = produce_test_messages(producer, source_topic)
file_consumer = ConsumerFactory().create_consumer(consumer_group,
source_topic_id,
output_directory,
False,
avro=False)
file_consumer.consume(10)
producer = ProducerFactory().create_producer(target_topic_id,
output_directory,
avro=False)
producer.produce()
# Check assertions:
assertion_check_directory = tmpdir_factory.mktemp(
"assertion_check_directory")
file_consumer = ConsumerFactory().create_consumer(
(consumer_group + "assertion_check"),
target_topic_id,
assertion_check_directory,
False,
avro=False)
file_consumer.consume(10)
consumed_messages = get_consumed_messages(assertion_check_directory, False)
assert produced_messages == consumed_messages
def main():
parser = argparse.ArgumentParser(
epilog="""Description:
Plays a video from a jpeg topic,
visualizes the head detection with an orage bounding box around a head
and writes demography data (gender & age) data above the heads.
Displays ('-d') or stores ('-o') the result of this demo in the kafka topic.
Required topics:
- <prefix>.cam.0.original.Image.jpg
- <prefix>.cam.0.dets.ObjectDetectionRecord.json
- <prefix>.cam.0.genders.GenderRecord.json
- <prefix>.cam.0.ages.AgeRecord.json
""",
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("broker",
help="The name of the kafka broker.",
type=str)
parser.add_argument("prefix",
help="Prefix of topics (base|skeleton).",
type=str)
parser.add_argument('-f', "--full_screen", action='store_true')
parser.add_argument('-d', "--display", action='store_true')
parser.add_argument('-o',
'--output',
help='write output image into kafka topic',
action='store_true')
args = parser.parse_args()
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
overlay = cv2.imread('resources/powered_by_white.png',
cv2.IMREAD_UNCHANGED)
image_topic = f"{args.prefix}.cam.0.original.Image.jpg"
detection_topic = f"{args.prefix}.cam.0.dets.ObjectDetectionRecord.json"
gender_topic = f"{args.prefix}.cam.0.genders.GenderRecord.json"
age_topic = f"{args.prefix}.cam.0.ages.AgeRecord.json"
output_topic_name = f"{args.prefix}.cam.0.demography.Image.jpg"
# handle full screen
window_name = "DEMO: Demography (gender & age)"
if args.full_screen:
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(args.broker, "detection", [
TopicInfo(image_topic),
TopicInfo(detection_topic),
TopicInfo(gender_topic),
TopicInfo(age_topic),
], 100, None, True)
i = 0
for msgs in consumer.getMessages():
for time, v in message_list_to_frame_structure(msgs).items():
img = v[args.prefix]["0"]["image"]
if type(img) == np.ndarray:
for head_detection in v[args.prefix]["0"]["head_detection"]:
object_detection_record = v[args.prefix]["0"][
"head_detection"][head_detection]["bounding_box"]
age_record = v[args.prefix]["0"]["head_detection"][
head_detection]["age"]
gender_record = v[args.prefix]["0"]["head_detection"][
head_detection]["gender"]
age = "" if age_record['age'] == {} else age_record['age']
gender = "" if gender_record[
'gender'] == {} else gender_record['gender']
# draw bounding_box
img = draw_nice_bounding_box(
img, object_detection_record["bounding_box"],
(10, 95, 255))
# write age and gender
img = draw_nice_text(
img,
str(gender) + " " + str(age),
object_detection_record["bounding_box"], (10, 95, 255))
# draw ultinous logo
img = draw_overlay(img, overlay, Position.BOTTOM_RIGHT)
# produce output topic
if args.output:
producer.produce(output_topic_name,
value=encode_image_to_message(img),
timestamp=time)
producer.poll(0)
if i % 100 == 0:
producer.flush()
i += 1
# display
if args.display:
cv2.imshow(window_name, img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
def main():
parser = argparse.ArgumentParser(
epilog="""Description:
Plays a video from a jpeg topic, visualizes the head detections and tracks, and pass detections.
Displays the result on screen ('-d') or stores result in kafka ('-o').
Required topics:
- <prefix>.cam.0.original.Image.jpg
- <prefix>.cam.0.dets.ObjectDetectionRecord.json
- <prefix>.cam.0.tracks.TrackChangeRecord.json
- <prefix>.cam.0.passdet.PassDetectionRecord.json
""",
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("broker",
help="The name of the kafka broker.",
type=str)
parser.add_argument("prefix",
help="Prefix of topics (base|skeleton).",
type=str)
parser.add_argument("config", help="Path to service config.", type=str)
parser.add_argument('-f', "--full_screen", action='store_true')
parser.add_argument('-d', "--display", action='store_true')
parser.add_argument('-v', "--video_file", action='store_true')
parser.add_argument('-o',
'--output',
help='write output image into kafka topic',
action='store_true')
args = parser.parse_args()
passdet_config_json = parse_config_data(args=args, parser=parser)
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
begin_flag = None
end_flag = EndFlag.NEVER
if args.video_file:
begin_flag = BeginFlag.BEGINNING
end_flag = EndFlag.END_OF_PARTITION
heartbeat_interval_ms = 1000
overlay = cv2.imread('resources/powered_by_white.png',
cv2.IMREAD_UNCHANGED)
passlines: Dict[str, PassLine] = {
pl["id"]: PassLine(next(pass_colors),
[(int(p["x"]), int(p["y"])) for p in pl["poly"]])
for pl in passdet_config_json["passLines"]
}
image_topic = f"{args.prefix}.cam.0.original.Image.jpg"
detection_topic = f"{args.prefix}.cam.0.dets.ObjectDetectionRecord.json"
track_topic = f"{args.prefix}.cam.0.tracks.TrackChangeRecord.json"
frameinfo_topic = f"{args.prefix}.cam.0.frameinfo.FrameInfoRecord.json"
passdet_topic = f"{args.prefix}.cam.0.passdet.PassDetectionRecord.json"
output_topic_name = f"{args.prefix}.cam.0.passdet.Image.jpg"
# Write notification if no message is received for this long
notification_delay_sec = 10
# handle full screen
window_name = "DEMO: Pass detection"
if args.full_screen:
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(
args.broker,
"detection", [
TopicInfo(image_topic),
TopicInfo(track_topic, drop=False),
TopicInfo(passdet_topic, drop=False),
TopicInfo(detection_topic),
TopicInfo(frameinfo_topic)
],
100,
None,
True,
begin_flag=begin_flag,
end_flag=end_flag,
heartbeat_interval_ms=heartbeat_interval_ms)
i = 0
scaling = 1.0
img_dimensions = (768, 1024)
last_image_ts = None
tracks: DefaultDict[Any, ColoredPolyLine] = defaultdict(
lambda: ColoredPolyLine(next(track_colors)))
for msgs in consumer.getMessages():
if not isinstance(msgs, HeartBeat):
for ts, v in message_list_to_frame_structure(msgs).items():
for track_key, track_val in v[
args.prefix]["0"]["track"].items():
if track_val["end_of_track"]:
if track_key in tracks:
del tracks[track_key]
continue
point = track_val["point"]["x"], track_val["point"]["y"]
tracks[track_key].add_point(point)
for pass_det in v[args.prefix]["0"]["passdet"].values():
if pass_det["type"] == "HEARTBEAT":
continue
elif pass_det["type"] == "END_OF_TRACK":
continue
elif pass_det["type"] == "PASS_CANDIDATE":
pass_id = pass_det["pass_candidate"]["pass"][
"pass_line_id"]
cross_dir = pass_det["pass_candidate"]["pass"][
"cross_dir"]
if pass_id in passlines:
passlines[pass_id].add_event(cross_dir)
elif pass_det["type"] == "PASS_REALIZED":
continue
img = v[args.prefix]["0"]["image"]
if type(img) != np.ndarray:
continue
last_image_ts = int(time.time())
# Set the image scale
img_dimensions = (img.shape[0], img.shape[1])
shape_orig = v[args.prefix]["0"]["head_detection"].pop(
"image", {})
if shape_orig:
scaling = img.shape[1] / shape_orig["frame_info"]["columns"]
# draw bounding_box
for head_detection in v[args.prefix]["0"]["head_detection"]:
object_detection_record = v[args.prefix]["0"][
"head_detection"][head_detection]["bounding_box"]
if object_detection_record["type"] == "PERSON_HEAD":
img = draw_nice_bounding_box(
canvas=img,
bounding_box=object_detection_record[
"bounding_box"],
color=(10, 95, 255),
scaling=scaling)
for t in tracks.values():
t.draw(img, scaling)
for idx, l in enumerate(passlines.values()):
l.draw(img, scaling)
cv2.putText(img,
"".join(l.events), (40, (idx + 1) * 50),
cv2.FONT_HERSHEY_COMPLEX,
2,
l.color,
5,
bottomLeftOrigin=True)
img = draw_overlay(canvas=img,
overlay=overlay,
position=Position.BOTTOM_RIGHT,
scale=scaling)
# produce output topic
if args.output:
producer.produce(output_topic_name,
value=encode_image_to_message(img),
timestamp=ts)
producer.poll(0)
if i % 100 == 0:
producer.flush()
i = 0
i += 1
# display
if args.display:
cv2.imshow(window_name, img)
# Write notification until the first message is received
# (output topic is not updated to ensure kafka timestamp consistency)
elif args.display and (
last_image_ts is None
or last_image_ts + notification_delay_sec < int(time.time())):
img = np.zeros((*img_dimensions, 3), np.uint8)
text = "Waiting for input Kafka topics to be populated. \n" \
"Please make sure that MGR and other necessary services are running."
img = draw_simple_text(canvas=img, text=text, color=(10, 95, 255))
cv2.imshow(window_name, img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
if args.video_file:
exit(130)
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
def main():
parser = init_parser()
args = parser.parse_args()
config_data = parse_config_data(args=args, parser=parser)
positive_areas = parse_areas(config_data, "positive_areas")
negative_areas = parse_areas(config_data, "negative_areas")
detection_types = parse_detection_types(config_data)
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
begin_flag = None
end_flag = EndFlag.NEVER
if args.video_file:
begin_flag = BeginFlag.BEGINNING
end_flag = EndFlag.END_OF_PARTITION
heartbeat_interval_ms = 1000
output_topic_name = f"{args.prefix}.cam.0.filtered_dets.Image.jpg"
# Write notification if no message is received for this long
notification_delay_sec = 10
# handle full screen
window_name = "DEMO: Filtered detection"
if args.full_screen:
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(
broker=args.broker,
groupid="detection",
topics_infos=[
TopicInfo(
f"{args.prefix}.cam.0.original.Image.jpg"), # image_topic
TopicInfo(
f"{args.prefix}.cam.0.filtered_dets.ObjectDetectionRecord.json"
), # filtered_detection_topic
TopicInfo(f"{args.prefix}.cam.0.dets.ObjectDetectionRecord.json"
) # detection_topic
],
latency_ms=100,
group_by_time=True,
begin_flag=begin_flag,
end_flag=end_flag,
heartbeat_interval_ms=heartbeat_interval_ms)
i = 0
scaling = 1.0
img_dimensions = (768, 1024)
last_image_ts = None
for msgs in consumer.getMessages():
if not isinstance(msgs, HeartBeat):
for ts, v in message_list_to_frame_structure(msgs).items():
frame_info = v[args.prefix]["0"]
img = frame_info["image"]
if type(img) != np.ndarray:
continue
last_image_ts = int(time.time())
# Set the image scale
img_dimensions = (img.shape[0], img.shape[1])
shape_orig = frame_info["head_detection"].pop("image", {})
if shape_orig:
scaling = img.shape[1] / shape_orig["frame_info"]["columns"]
# draw bounding_box
for head_detection in frame_info["head_detection"]:
img = draw_bounding_box(
object_detection_record=frame_info["head_detection"]
[head_detection]["bounding_box"],
detection_types=detection_types,
img=img,
scaling=scaling,
color=COLOR_GREY)
for head_detection in frame_info["filtered_head_detection"]:
img = draw_bounding_box(object_detection_record=frame_info[
"filtered_head_detection"][head_detection]
["filtered_bounding_box"],
detection_types=detection_types,
img=img,
scaling=scaling,
color=COLOR_ORANGE)
draw_areas(areas=positive_areas, img=img, color=COLOR_GREEN)
draw_areas(areas=negative_areas, img=img, color=COLOR_RED)
draw_ultinous_logo(canvas=img, scale=scaling)
# produce output topic
if args.output:
producer.produce(output_topic_name,
value=encode_image_to_message(img),
timestamp=ts)
producer.poll(0)
if i % 100 == 0:
producer.flush()
i = 0
i += 1
# display
if args.display:
cv2.imshow(window_name, img)
# Write notification until the first message is received
# (output topic is not updated to ensure kafka timestamp consistency)
elif args.display and (
last_image_ts is None
or last_image_ts + notification_delay_sec < int(time.time())):
img = np.zeros((*img_dimensions, 3), np.uint8)
text = "Waiting for input Kafka topics to be populated. \n" \
"Please make sure that MGR and other necessary services are running."
img = draw_simple_text(canvas=img, text=text, color=(10, 95, 255))
cv2.imshow(window_name, img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
if args.video_file:
exit(130)
else:
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
def main(args):
serial = args.serial
num_messages = args.num_messages
brokers = args.brokers
group_id = args.group_id
input_topic = args.input_topic
input_partition = args.input_partition
output_topic = args.output_topic
if serial:
print("Running in SERIAL mode")
print(
"The input producer will wait for the reply of the transactor before producing the next message."
)
else:
print("Running in PARALLEL mode")
print(
"The input producer will produce all messages in parallel (at once) after the first message."
)
tr_args = [
sys.executable,
os.path.join(HERE, "eos-transactions.py"),
"-b",
brokers,
"-g",
group_id + "-tr",
"-t",
input_topic,
"-p",
str(input_partition),
"-o",
output_topic,
]
output_consumer = Consumer({
"bootstrap.servers": brokers,
"group.id": group_id + "-pr",
"auto.offset.reset": "earliest",
"enable.auto.commit": True,
"enable.partition.eof": False,
})
output_consumer.subscribe([output_topic])
input_producer = Producer({
'bootstrap.servers': brokers,
})
try:
with tempfile.NamedTemporaryFile(mode='w+') as f:
tr_proc = subprocess.Popen(tr_args,
stderr=subprocess.STDOUT,
stdout=f,
cwd=HERE,
close_fds=True)
try:
time.sleep(1)
assert tr_proc.poll() is None
tx = 0
for i in range(num_messages):
input_producer.produce(input_topic,
key=b"xy",
value=str(tx).encode("ascii"))
tx += 1
assert input_producer.flush(10) == 0
while serial or tx <= 1:
msg = output_consumer.poll(1.0)
if msg is None:
continue
assert msg.error() is None
if tx == 1:
t_start = time.time()
break
if not serial:
for _ in range(num_messages - 1):
msg = output_consumer.poll(1.0)
if msg is None:
continue
assert msg.error() is None
print("Processing took {}".format(time.time() - t_start))
finally:
if tr_proc.poll() is None:
tr_proc.terminate()
tr_proc.wait()
f.seek(0)
eos_out = f.read()
finally:
output_consumer.close() # commit offsets
i = 0
c = False
send_offset_logs = defaultdict(list)
send_offset_times = []
for line in eos_out.split("\n"):
if line.startswith(":DEMO:START "):
c = True
if c:
send_offset_logs[i].append(line)
if line.startswith(":DEMO:END "):
send_offset_times.append(float(line.rpartition(" ")[-1]))
c = False
i += 1
print("\nSend offset times:", send_offset_times)
print("Send offset times average:",
sum(send_offset_times) / len(send_offset_times))
print("\nRelevant log snippet from the middle:")
print("\n".join(send_offset_logs[int(i / 2)]))
print("\nFull output of the transactor:")
print(eos_out)
def main():
parser = argparse.ArgumentParser(
epilog="""Description:
Basic Reidentification demo using two cameras:
Camera0 for object registration and Camera1 for reidentification.
Plays a video from a jpeg topic,
visualizes head detection with an orange bounding box around a head
and writes the dwell time and ID (derived from the reid MS ID) above the heads.
Displays ('-d') or stores ('-o') the result of this demo in kafka topics.
Required topics:
- <prefix>.cam.0.original.Image.jpg
- <prefix>.cam.0.dets.ObjectDetectionRecord.json
- <prefix>.cam.1.original.Image.jpg
- <prefix>.cam.1.dets.ObjectDetectionRecord.json
- <prefix>.cam.1.reids.ReidRecord.json
""",
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("broker",
help="The name of the kafka broker.",
type=str)
parser.add_argument("prefix",
help="Prefix of topics (base|skeleton).",
type=str)
parser.add_argument('-f', "--full_screen", action='store_true')
parser.add_argument('-d', "--display", action='store_true')
parser.add_argument('-o',
'--output',
help='write output image into kafka topic',
action='store_true')
args = parser.parse_args()
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
overlay = cv2.imread('resources/powered_by_white.png',
cv2.IMREAD_UNCHANGED)
image_reg_topic = f"{args.prefix}.cam.{REG_CAMERA_ID}.original.Image.jpg"
image_reid_topic = f"{args.prefix}.cam.{REID_CAMERA_ID}.original.Image.jpg"
detection_reg_topic = f"{args.prefix}.cam.{REG_CAMERA_ID}.dets.ObjectDetectionRecord.json"
detection_reid_topic = f"{args.prefix}.cam.{REID_CAMERA_ID}.dets.ObjectDetectionRecord.json"
reid_topic = f"{args.prefix}.cam.{REID_TOPIC_ID}.reids.ReidRecord.json"
output_reg_topic_name = f"{args.prefix}.cam.{REG_CAMERA_ID}.reids.Image.jpg"
output_reid_topic_name = f"{args.prefix}.cam.{REID_CAMERA_ID}.reids.Image.jpg"
# handle full screen
window_name = TITLE
if args.full_screen:
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(
broker=args.broker,
groupid="detection",
topics_infos=[
TopicInfo(image_reg_topic),
TopicInfo(image_reid_topic),
TopicInfo(detection_reg_topic),
TopicInfo(detection_reid_topic),
TopicInfo(reid_topic, drop=False),
],
latency_ms=500,
commit_interval_sec=None,
group_by_time=True)
registrations: Dict[str, int] = {}
i = 0
inner_id = 0
for msgs in consumer.getMessages():
for time, v in message_list_to_frame_structure(msgs).items():
message = v.get(args.prefix, {})
# Register the recognized faces
reid_records = message[REID_TOPIC_ID].get("reid", {})
for reid_key, reid_record in reid_records.items():
record_type = reid_record["type"]
# Get the stored registration key
registration_key = reid_record["reg_refs"][0]["subject"]["key"]
if record_type == "REG" and registration_key not in registrations:
inner_id += 1
registrations[registration_key] = inner_id
for topic_key, topic_message in message.items():
img = topic_message.get("image", {})
if not isinstance(img, np.ndarray):
continue
# Process detections
head_detections = topic_message.get("head_detection", {})
for detection_key, detection_record in head_detections.items():
object_detection_record = detection_record["bounding_box"]
color = COLOR_GREY
key_to_display = ""
# Reidentification received
reid_record = reid_records.get(detection_key)
if reid_record and reid_record["type"] == "REID":
reid_key = reid_record["reg_refs"][0]["subject"][
"key"] # We only use the first identified face now
registered_id = registrations.get(reid_key)
if registered_id:
color = COLOR_ORANGE
dwell_time = time - int(reid_key.split('_')[0])
key_to_display = f"id: {registered_id}; dwell time: {dwell_time}ms"
# draw text above bounding box
img = draw_nice_text(
canvas=img,
text=key_to_display,
bounding_box=object_detection_record["bounding_box"],
color=color)
# draw bounding_box
img = draw_nice_bounding_box(
canvas=img,
bounding_box=object_detection_record["bounding_box"],
color=color)
# draw ultinous logo
img = draw_overlay(canvas=img,
overlay=overlay,
position=Position.BOTTOM_RIGHT)
# produce output topic
if args.output:
out_topic = output_reg_topic_name if topic_key is REG_CAMERA_ID else output_reid_topic_name
producer.produce(out_topic,
value=encode_image_to_message(img),
timestamp=time)
producer.poll(0)
if i % 100 == 0:
producer.flush()
i += 1
# display #
if args.display:
cv2.imshow(window_name, img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
def main():
parser = argparse.ArgumentParser(
epilog="""Description:
Plays a video from a jpeg topic,
visualizes the head detection with a bounding box around a head.
The boundig box is grey when mask detection did not run;
it is green when a mask is detected;
it is orange and 'NO MASK' is written above the head when no mask is detected.
Displays ('-d') or stores ('-o') the result of this demo in the kafka topic.
Required topics:
- <prefix>.cam.0.original.Image.jpg
- <prefix>.cam.0.dets.ObjectDetectionRecord.json
- <prefix>.cam.0.masks.FaceMaskRecord.json
""",
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("broker",
help="The name of the kafka broker.",
type=str)
parser.add_argument("prefix",
help="Prefix of topics (base|skeleton).",
type=str)
parser.add_argument('-f', "--full_screen", action='store_true')
parser.add_argument('-d', "--display", action='store_true')
parser.add_argument('-v', "--video_file", action='store_true')
parser.add_argument('-o',
'--output',
help='write output image into kafka topic',
action='store_true')
args = parser.parse_args()
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
begin_flag = None
end_flag = EndFlag.NEVER
if args.video_file:
begin_flag = BeginFlag.BEGINNING
end_flag = EndFlag.END_OF_PARTITION
heartbeat_interval_ms = 1000
overlay = cv2.imread('resources/powered_by_white.png',
cv2.IMREAD_UNCHANGED)
image_topic = f"{args.prefix}.cam.0.original.Image.jpg"
detection_topic = f"{args.prefix}.cam.0.dets.ObjectDetectionRecord.json"
mask_topic = f"{args.prefix}.cam.0.masks.FaceMaskRecord.json"
output_topic_name = f"{args.prefix}.cam.0.face_mask.Image.jpg"
frameinfo_topic = f"{args.prefix}.cam.0.frameinfo.FrameInfoRecord.json"
# Write notification if no message is received for this long
notification_delay_sec = 10
# handle full screen
window_name = "DEMO: Face Mask"
if args.full_screen:
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(
args.broker,
"detection", [
TopicInfo(image_topic),
TopicInfo(detection_topic),
TopicInfo(mask_topic),
TopicInfo(frameinfo_topic)
],
100,
None,
True,
begin_flag=begin_flag,
end_flag=end_flag,
heartbeat_interval_ms=heartbeat_interval_ms)
i = 0
scaling = 1.0
img_dimensions = (768, 1024)
last_image_ts = None
for msgs in consumer.getMessages():
if not isinstance(msgs, HeartBeat):
for ts, v in message_list_to_frame_structure(msgs).items():
img = v[args.prefix]["0"]["image"]
if type(img) != np.ndarray:
continue
last_image_ts = int(time.time())
# Set the image scale
img_dimensions = (img.shape[0], img.shape[1])
shape_orig = v[args.prefix]["0"]["head_detection"].pop(
"image", {})
if shape_orig:
scaling = img.shape[1] / shape_orig["frame_info"]["columns"]
for head_detection in v[args.prefix]["0"]["head_detection"]:
object_detection_record = v[args.prefix]["0"][
"head_detection"][head_detection]["bounding_box"]
if object_detection_record["type"] != "PERSON_HEAD":
continue
mask_record = v[args.prefix]["0"]["head_detection"][
head_detection]["face_mask"]
mask_text = ""
if not mask_record:
color = COLOR_DARK_GREY
elif mask_record["has_mask"]:
color = COLOR_GREEN
else:
mask_text = "NO MASK"
color = COLOR_ORANGE
# draw bounding_box
img = draw_nice_bounding_box(
canvas=img,
bounding_box=object_detection_record["bounding_box"],
color=color,
scaling=scaling)
# write age and gender
img = draw_nice_text(
img,
mask_text,
object_detection_record["bounding_box"],
color,
scale=scaling)
# draw ultinous logo
img = draw_overlay(canvas=img,
overlay=overlay,
position=Position.BOTTOM_RIGHT,
scale=scaling)
# produce output topic
if args.output:
producer.produce(output_topic_name,
value=encode_image_to_message(img),
timestamp=ts)
producer.poll(0)
if i % 100 == 0:
producer.flush()
i += 1
# display
if args.display:
cv2.imshow(window_name, img)
# Write notification until the first message is received
# (output topic is not updated to ensure kafka timestamp consistency)
elif args.display and (
last_image_ts is None
or last_image_ts + notification_delay_sec < int(time.time())):
img = np.zeros((*img_dimensions, 3), np.uint8)
text = "Waiting for input Kafka topics to be populated. \n" \
"Please make sure that MGR and other necessary services are running."
img = draw_simple_text(canvas=img, text=text, color=(10, 95, 255))
cv2.imshow(window_name, img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
if args.video_file:
exit(130)
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
def main():
parser = argparse.ArgumentParser(
epilog="""Description:
Plays a video from a jpeg topic and visualizes the head detection with an orange bounding box around a head.
Displays ('-d') or stores ('-o') the result of this demo in the kafka topic.
Required topics:
- <prefix>.cam.0.original.Image.jpg
- <prefix>.cam.0.dets.ObjectDetectionRecord.json
""",
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("broker",
help="The name of the kafka broker.",
type=str)
parser.add_argument("prefix",
help="Prefix of topics (base|skeleton).",
type=str)
parser.add_argument('-f', "--full_screen", action='store_true')
parser.add_argument('-d', "--display", action='store_true')
parser.add_argument('-v', "--video_file", action='store_true')
parser.add_argument('-o',
'--output',
help='write output image into kafka topic',
action='store_true')
args = parser.parse_args()
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
begin_flag = None
end_flag = None
if args.video_file:
begin_flag = BeginFlag.BEGINNING
end_flag = EndFlag.END_OF_PARTITION
overlay = cv2.imread('resources/powered_by_white.png',
cv2.IMREAD_UNCHANGED)
image_topic = f"{args.prefix}.cam.0.original.Image.jpg"
detection_topic = f"{args.prefix}.cam.0.dets.ObjectDetectionRecord.json"
frameinfo_topic = f"{args.prefix}.cam.0.frameinfo.FrameInfoRecord.json"
output_topic_name = f"{args.prefix}.cam.0.head_detection.Image.jpg"
# handle full screen
window_name = "DEMO: Head detection"
if args.full_screen:
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(args.broker,
"detection", [
TopicInfo(image_topic),
TopicInfo(detection_topic),
TopicInfo(frameinfo_topic)
],
100,
None,
True,
begin_flag=begin_flag,
end_flag=end_flag)
i = 0
scaling = 1.0
for msgs in consumer.getMessages():
for time, v in message_list_to_frame_structure(msgs).items():
img = v[args.prefix]["0"]["image"]
if type(img) == np.ndarray:
# Set the image scale
shape_orig = v[args.prefix]["0"]["head_detection"].pop(
"image", {})
if shape_orig:
scaling = img.shape[1] / shape_orig["frame_info"]["columns"]
# draw bounding_box
for head_detection in v[args.prefix]["0"]["head_detection"]:
object_detection_record = v[args.prefix]["0"][
"head_detection"][head_detection]["bounding_box"]
if object_detection_record["type"] == "PERSON_HEAD":
img = draw_nice_bounding_box(
canvas=img,
bounding_box=object_detection_record[
"bounding_box"],
color=(10, 95, 255),
scaling=scaling)
# draw ultinous logo
img = draw_overlay(canvas=img,
overlay=overlay,
position=Position.BOTTOM_RIGHT,
scale=scaling)
# produce output topic
if args.output:
producer.produce(output_topic_name,
value=encode_image_to_message(img),
timestamp=time)
producer.poll(0)
if i % 100 == 0:
producer.flush()
i = 0
i += 1
# display
if args.display:
cv2.imshow(window_name, img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
if args.video_file:
exit(130)
else:
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
def main():
parser = init_parser()
args = parser.parse_args()
config_data = parse_config_data(args=args, parser=parser)
positive_areas = parse_areas(config_data, "positive_areas")
negative_areas = parse_areas(config_data, "negative_areas")
detection_types = parse_detection_types(config_data)
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
begin_flag = None
end_flag = None
if args.video_file:
begin_flag = BeginFlag.BEGINNING
end_flag = EndFlag.END_OF_PARTITION
output_topic_name = f"{args.prefix}.cam.0.filtered_dets.Image.jpg"
# handle full screen
window_name = "DEMO: Filtered detection"
if args.full_screen:
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(
broker=args.broker,
groupid="detection",
topics_infos=[
TopicInfo(
f"{args.prefix}.cam.0.original.Image.jpg"), # image_topic
TopicInfo(
f"{args.prefix}.cam.0.filtered_dets.ObjectDetectionRecord.json"
), # filtered_detection_topic
TopicInfo(f"{args.prefix}.cam.0.dets.ObjectDetectionRecord.json"
) # detection_topic
],
latency_ms=100,
group_by_time=True,
begin_flag=begin_flag,
end_flag=end_flag)
i = 0
scaling = 1.0
for msgs in consumer.getMessages():
for time, v in message_list_to_frame_structure(msgs).items():
frame_info = v[args.prefix]["0"]
img = frame_info["image"]
if type(img) == np.ndarray:
# Set the image scale
shape_orig = frame_info["head_detection"].pop("image", {})
if shape_orig:
scaling = img.shape[1] / shape_orig["frame_info"]["columns"]
# draw bounding_box
for head_detection in frame_info["head_detection"]:
img = draw_bounding_box(
object_detection_record=frame_info["head_detection"]
[head_detection]["bounding_box"],
detection_types=detection_types,
img=img,
scaling=scaling,
color=COLOR_GREY)
for head_detection in frame_info["filtered_head_detection"]:
img = draw_bounding_box(object_detection_record=frame_info[
"filtered_head_detection"][head_detection]
["filtered_bounding_box"],
detection_types=detection_types,
img=img,
scaling=scaling,
color=COLOR_ORANGE)
draw_areas(areas=positive_areas, img=img, color=COLOR_GREEN)
draw_areas(areas=negative_areas, img=img, color=COLOR_RED)
draw_ultinous_logo(canvas=img, scale=scaling)
# produce output topic
if args.output:
producer.produce(output_topic_name,
value=encode_image_to_message(img),
timestamp=time)
producer.poll(0)
if i % 100 == 0:
producer.flush()
i = 0
i += 1
# display
if args.display:
cv2.imshow(window_name, img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
if args.video_file:
exit(130)
else:
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
def main():
parser = argparse.ArgumentParser(
epilog="""Description:
Plays a video from a jpeg topic, visualizes the head detections and tracks, and pass detections.
Displays the result on screen ('-d') or stores result in kafka ('-o').
Required topics:
- <prefix>.cam.0.original.Image.jpg
- <prefix>.cam.0.dets.ObjectDetectionRecord.json
- <prefix>.cam.0.tracks.TrackChangeRecord.json
- <prefix>.cam.0.passdet.PassDetectionRecord.json
""",
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("broker",
help="The name of the kafka broker.",
type=str)
parser.add_argument("prefix",
help="Prefix of topics (base|skeleton).",
type=str)
parser.add_argument("config", help="Path to service config.", type=str)
parser.add_argument('-f', "--full_screen", action='store_true')
parser.add_argument('-d', "--display", action='store_true')
parser.add_argument('-o',
'--output',
help='write output image into kafka topic',
action='store_true')
args = parser.parse_args()
config_file = Path(args.config)
if not config_file.is_file():
parser.error(f"{args.config} does not exist.")
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
with config_file.open() as f:
try:
passdet_config_json = json.loads(
javaproperties.load(f)
["ultinous.service.kafka.passdet.config"])
except KeyError:
parser.error(
"Missing property: ultinous.service.kafka.passdet.config")
except JSONDecodeError as e:
parser.error(f"Error parsing {e}")
overlay = cv2.imread('resources/powered_by_white.png',
cv2.IMREAD_UNCHANGED)
passlines: Dict[str, PassLine] = {
pl["id"]: PassLine(next(pass_colors),
[(int(p["x"]), int(p["y"])) for p in pl["poly"]])
for pl in passdet_config_json["passLines"]
}
image_topic = f"{args.prefix}.cam.0.original.Image.jpg"
detection_topic = f"{args.prefix}.cam.0.dets.ObjectDetectionRecord.json"
track_topic = f"{args.prefix}.cam.0.tracks.TrackChangeRecord.json"
passdet_topic = f"{args.prefix}.cam.0.passdet.PassDetectionRecord.json"
output_topic_name = f"{args.prefix}.cam.0.passdet.Image.jpg"
# handle full screen
window_name = "DEMO: Pass detection"
if args.full_screen:
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(args.broker, "detection", [
TopicInfo(image_topic),
TopicInfo(track_topic, drop=False),
TopicInfo(passdet_topic, drop=False),
TopicInfo(detection_topic)
], 100, None, True)
i = 0
tracks: DefaultDict[Any, ColoredPolyLine] = defaultdict(
lambda: ColoredPolyLine(next(track_colors)))
for msgs in consumer.getMessages():
for time, v in message_list_to_frame_structure(msgs).items():
for track_key, track_val in v[args.prefix]["0"]["track"].items():
if track_val["end_of_track"]:
if track_key in tracks:
del tracks[track_key]
continue
point = track_val["point"]["x"], track_val["point"]["y"]
tracks[track_key].add_point(point)
for pass_det in v[args.prefix]["0"]["passdet"].values():
if pass_det["type"] == "HEARTBEAT":
continue
elif pass_det["type"] == "END_OF_TRACK":
continue
elif pass_det["type"] == "PASS_CANDIDATE":
pass_id = pass_det["pass_candidate"]["pass"][
"pass_line_id"]
cross_dir = pass_det["pass_candidate"]["pass"]["cross_dir"]
if pass_id in passlines:
passlines[pass_id].add_event(cross_dir)
elif pass_det["type"] == "PASS_REALIZED":
continue
img = v[args.prefix]["0"]["image"]
if type(img) == np.ndarray:
# draw bounding_box
for head_detection in v[args.prefix]["0"]["head_detection"]:
object_detection_record = v[args.prefix]["0"][
"head_detection"][head_detection]["bounding_box"]
if object_detection_record["type"] == "PERSON_HEAD":
img = draw_nice_bounding_box(
img, object_detection_record["bounding_box"],
(10, 95, 255))
for t in tracks.values():
t.draw(img)
for idx, l in enumerate(passlines.values()):
l.draw(img)
cv2.putText(img,
"".join(l.events), (40, (idx + 1) * 50),
cv2.FONT_HERSHEY_COMPLEX,
2,
l.color,
5,
bottomLeftOrigin=True)
img = draw_overlay(img, overlay, Position.BOTTOM_RIGHT)
# produce output topic
if args.output:
producer.produce(output_topic_name,
value=encode_image_to_message(img),
timestamp=time)
producer.poll(0)
if i % 100 == 0:
producer.flush()
i = 0
i += 1
# display
if args.display:
cv2.imshow(window_name, img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
def main():
parser = argparse.ArgumentParser(
epilog="""Description:
Plays video from a jpeg topic, visualizes main points of a human skeleton linked with colorful lines.
Displays the result on screen ('-d') or stores in a kafka topic with '-o' parameter.
Required topics:
- <prefix>.cam.0.original.Image.jpg
- <prefix>.cam.0.skeletons.SkeletonRecord.json
""",
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("broker",
help="The name of the kafka broker.",
type=str)
parser.add_argument("prefix",
help="Prefix of topics (base|skeleton).",
type=str)
parser.add_argument('-f', "--full_screen", action='store_true')
parser.add_argument('-d', "--display", action='store_true')
parser.add_argument('-v', "--video_file", action='store_true')
parser.add_argument('-o',
'--output',
help='write output image into kafka topic',
action='store_true')
args = parser.parse_args()
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
begin_flag = None
end_flag = None
if args.video_file:
begin_flag = BeginFlag.BEGINNING
end_flag = EndFlag.END_OF_PARTITION
overlay = cv2.imread('resources/powered_by_white.png',
cv2.IMREAD_UNCHANGED)
img_topic = f"{args.prefix}.cam.0.original.Image.jpg"
skeleton_topic = f"{args.prefix}.cam.0.skeletons.SkeletonRecord.json"
output_topic_name = f"{args.prefix}.cam.0.skeleton.Image.jpg"
# handle full screen
window_name = "DEMO: Human skeleton"
if args.full_screen:
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(
args.broker,
"skeleton",
[TopicInfo(img_topic), TopicInfo(skeleton_topic)],
100,
None,
True,
begin_flag=begin_flag,
end_flag=end_flag)
i = 0
for msgs in consumer.getMessages():
for time, v in message_list_to_frame_structure(msgs).items():
img = v[args.prefix]["0"]["image"]
if type(img) == np.ndarray:
# draw skeletons
for skeleton_id, skeleton in v[
args.prefix]["0"]["skeleton"].items():
img = draw_skeleton_with_background(
img, skeleton["points"])
# draw ultinous logo
img = draw_overlay(img, overlay, Position.BOTTOM_RIGHT)
# produce output topic
if args.output:
producer.produce(output_topic_name,
value=encode_image_to_message(img),
timestamp=time)
producer.poll(0)
if i % 100 == 0:
producer.flush()
i += 1
# display
if args.display:
cv2.imshow(window_name, img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
if args.video_file:
exit(130)
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
class Kafka(object):
def __init__(self, target_key) -> None:
super().__init__()
self.address = _address_for_key(target_key)
kafka_config = {
'bootstrap.servers': self.address,
'group.id': "up9-test-group",
'enable.auto.commit':
'false' # important for passive observing
}
if "ssl://" in self.address.lower():
kafka_config['security.protocol'] = 'SSL'
self.consumer = Consumer(kafka_config)
self.producer = Producer(kafka_config)
self.watching_topics = []
self.consumer.list_topics(timeout=5) # to check for connectivity
def watch_topics(self, topics: list):
def my_on_assign(consumer, partitions):
logging.debug("On assign: %r", partitions)
consumer.assign(partitions)
for partition in partitions:
low, high = consumer.get_watermark_offsets(partition)
partition.offset = high
logging.debug("Setting offset: %r", partition)
consumer.seek(partition)
self.watching_topics.extend(topics)
self.consumer.subscribe(topics, on_assign=my_on_assign)
self.consumer.poll(0.01) # to trigger partition assignments
def get_watched_messages(self, interval=0.0, predicate=lambda x: True):
logging.debug(
"Checking messages that appeared on kafka topics: %r",
self.watching_topics)
res = []
start = time.time()
while True:
msg = self.consumer.poll(interval)
if msg is None or time.time() - start > interval:
break # done reading
if msg.error():
raise KafkaException("kafka consumer error: {}".format(
msg.error()))
logging.debug(
"Potential message: %r",
(msg.partition(), msg.key(), msg.headers(), msg.value()))
if predicate(msg):
res.append(msg)
# TODO: consumer.close()
return res
def assert_seen_message(self, resp, delay=0, predicate=lambda x: True):
@recorder.assertion_decorator
def assert_seen_kafka_message(resp, topics, delay):
messages = self.get_watched_messages(delay, predicate)
messages = [(m.topic(), m.key(), m.value(), m.headers())
for m in messages]
if not messages:
raise AssertionError("No messages on Kafka topic %r" %
topics)
else:
logging.info("Validated the messages have appeared: %s",
messages)
return messages
return assert_seen_kafka_message(resp, self.watching_topics, delay)
def put(self, topic, data=None, json=None, headers=None):
# TODO: parse key out of URL
if topic.startswith('/'):
topic = topic[1:]
if data is None and json is not None:
data = json_lib.dumps(json)
with apiritif.transaction('kafka://[' + self.address + ']/' +
topic):
logging.info("Sending message to Kafka topic %r: %r", topic,
data)
self.producer.produce(
topic, data, headers=[] if headers is None else headers)
self.producer.poll(0)
self.producer.flush()
wrapped_req = self._make_request(
'PUT',
'kafka://' + self.address.split(',')[0] + '/' + topic,
data)
wrapped_response = self._make_response(wrapped_req)
recorder.record_http_request('PUT', self.address, wrapped_req,
wrapped_response,
_context.session)
return wrapped_response
def _make_request(self, method, url, request):
req = requests.Request(method, url=url, data=request)
prepared = req.prepare()
_context.grpc_mapping[id(request)] = prepared
return prepared
def _make_response(self, wrapped_req):
resp = requests.Response()
resp.status_code = 202
resp.request = wrapped_req
resp._request = wrapped_req
resp.msg = 'Accepted'
resp.raw = io.BytesIO()
return resp
def main():
parser = argparse.ArgumentParser(
epilog="""Description:
Plays a video from a jpeg topic,
visualizes head detection with an orage bounding box around a head
and writes the IDs given by reid MS above the heads.
Displays ('-d') or stores ('-o') the result of this demo in the kafka topic.
Required topics:
- <prefix>.cam.0.original.Image.jpg
- <prefix>.cam.0.dets.ObjectDetectionRecord.json
- <prefix>.cam.99.reids.ReidRecord.json
""",
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("broker",
help="The name of the kafka broker.",
type=str)
parser.add_argument("prefix",
help="Prefix of topics (base|skeleton).",
type=str)
parser.add_argument('-f', "--full_screen", action='store_true')
parser.add_argument('-d', "--display", action='store_true')
parser.add_argument('-v', "--video_file", action='store_true')
parser.add_argument('-o',
'--output',
help='write output image into kafka topic',
action='store_true')
args = parser.parse_args()
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
begin_flag = None
end_flag = None
if args.video_file:
begin_flag = BeginFlag.BEGINNING
end_flag = EndFlag.END_OF_PARTITION
overlay = cv2.imread('resources/powered_by_white.png',
cv2.IMREAD_UNCHANGED)
image_topic = f"{args.prefix}.cam.0.original.Image.jpg"
detection_topic = f"{args.prefix}.cam.0.dets.ObjectDetectionRecord.json"
reid_topic = f"{args.prefix}.cam.{REID_TOPIC_ID}.reids.ReidRecord.json"
output_topic_name = f"{args.prefix}.cam.0.reidentification.Image.jpg"
frameinfo_topic = f"{args.prefix}.cam.0.frameinfo.FrameInfoRecord.json"
# handle full screen
window_name = TITLE
if args.full_screen:
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(
args.broker,
"detection",
[
TopicInfo(image_topic),
TopicInfo(detection_topic),
TopicInfo(reid_topic),
TopicInfo(frameinfo_topic)
],
500,
None,
True,
begin_flag=begin_flag,
end_flag=end_flag,
)
i = 0
stored_ids = {}
scaling = 1.0
for msgs in consumer.getMessages():
for time, v in message_list_to_frame_structure(msgs).items():
message = v.get(args.prefix, {})
reid_records = message[REID_TOPIC_ID].get("reid", {})
img = message["0"].get("image", {})
if type(img) == np.ndarray:
head_detections = message["0"].get("head_detection", {})
# Set the image scale
shape_orig = head_detections.pop("image", {})
if shape_orig:
scaling = img.shape[1] / shape_orig["frame_info"]["columns"]
# Processing detections
for detection_key, detection_record in head_detections.items():
object_detection_record = detection_record["bounding_box"]
color = COLOR_GREY
reid_records_for_det = reid_records.get(detection_key, ())
for reid_record in filter(lambda r: "reid_event" in r,
reid_records_for_det):
color = COLOR_ORANGE
reid_key = reid_record["reid_event"]["match_list"][0][
"id"]["first_detection_key"]
key_to_display = stored_ids.get(reid_key, None)
if key_to_display is None:
key_to_display = len(stored_ids) + 1
stored_ids[reid_key] = key_to_display
# user id
img = draw_nice_text(
canvas=img,
text=str(key_to_display),
bounding_box=object_detection_record[
"bounding_box"],
color=color,
scale=scaling)
# draw bounding_box
img = draw_nice_bounding_box(
canvas=img,
bounding_box=object_detection_record["bounding_box"],
color=color,
scaling=scaling)
# draw ultinous logo
img = draw_overlay(canvas=img,
overlay=overlay,
position=Position.BOTTOM_RIGHT,
scale=scaling)
# produce output topic
if args.output:
producer.produce(output_topic_name,
value=encode_image_to_message(img),
timestamp=time)
producer.poll(0)
if i % 100 == 0:
producer.flush()
i += 1
# display
if args.display:
cv2.imshow(window_name, img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
if args.video_file:
exit(130)
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
class KafkaWriterConfiguration:
def __init__(self, configuration):
self.broker = configuration['broker']
self.topics = []
for topic_name, topic_configuration in configuration['topics'].items():
self.topics.append(
KafkaConfigurationTopic(topic_name,
topic_configuration['replication'],
topic_configuration['partitions'],
topic_configuration['recreate']))
self.producer = None
self.producer_configuration = configuration['producer'][
'configuration']
def create_or_recreate_topics(self, admin=None):
topics_to_recreate = list(
filter(lambda topic_spec: topic_spec.recreate, self.topics))
if topics_to_recreate:
logger.info('Got %s topics to recreate', topics_to_recreate)
if not admin:
admin = AdminClient({'bootstrap.servers': self.broker})
topics_to_delete = list(
map(lambda topic_spec: topic_spec.name, topics_to_recreate))
deleted_topics_result = admin.delete_topics(topics_to_delete,
operation_timeout=15,
request_timeout=45)
for deleted_topic, future in deleted_topics_result.items():
try:
future.result()
except Exception as e:
# TODO: transform to WARN
print('An error occurred on deleting {}: {}'.format(
deleted_topic, e))
# Apparently there is a concurrent issue with topics create or recreate operations
# https://github.com/confluentinc/confluent-kafka-python/issues/524
# That's why I will try to create the topic 5 times with a kind of exponential backoff
topics_to_create = {
topic_spec.name: topic_spec.to_new_topic()
for topic_spec in topics_to_recreate
}
creation_errors = []
for retry in range(0, 5):
if topics_to_create:
futures = admin.create_topics(list(
topics_to_create.values()),
operation_timeout=15,
request_timeout=45)
for topic, f in futures.items():
try:
f.result()
print("Topic {} created".format(topic))
del topics_to_create[topic]
except Exception as e:
logger.warning(
"Topic %s was not created - retrying. Got exception %s",
topic,
e,
exc_info=1)
if retry == 4:
creation_errors.append((topic, e))
time.sleep(retry * 5)
assert not creation_errors, "Some topics were not correctly created {}".format(
list(
map(
lambda topic_and_error: "{}: {}".format(
topic_and_error[0], topic_and_error[1]),
creation_errors)))
def send_message(self, topic_name, key, message):
if not self.producer:
# check https://github.com/edenhill/librdkafka/blob/master/CONFIGURATION.md for more information
producer_config = {'bootstrap.servers': self.broker}
producer_config.update(self.producer_configuration)
self.producer = Producer(producer_config)
"""
Traceback (most recent call last):
File "/home/bartosz/workspace/data-generator/examples/kafka/generate_dataset_to_kafka.py", line 51, in <module>
configuration.send_message(output_topic_name, action)
File "/home/bartosz/workspace/data-generator/data_generator/sink/kafka_writer.py", line 60, in send_message
self.producer.produce(topic_name, value=bytes(message, encoding='utf-8'))
BufferError: Local: Queue full
The below code is the workaround for the above problem, found here:
`Confluent Kafka-Python issue 104 <https://github.com/confluentinc/confluent-kafka-python/issues/104>`
"""
def delivery_callback(error, result):
if error:
logger.error("Record was not correctly delivered: %s", error)
# if the key is missing, it will fail with "TypeError: encoding without a string argument" error
# replace the missing key by some dummy value
message_key = key if key else 'empty'
try:
self.producer.produce(topic=topic_name,
key=bytes(message_key, encoding='utf-8'),
value=bytes(message, encoding='utf-8'),
on_delivery=delivery_callback)
except BufferError:
self.producer.flush()
self.producer.produce(topic=topic_name,
key=bytes(message_key, encoding='utf-8'),
value=bytes(message, encoding='utf-8'),
on_delivery=delivery_callback)
def __repr__(self):
return 'KafkaWriterConfiguration (broker={}) (topics={})'.format(
self.broker, self.topics)
class ModelInference(object):
def __init__(self,
my_id=1,
bootstrap_servers='',
list_of_partitions=[],
request_topic='',
inference_topic='',
group_id='my_grp'):
""" Constructor
:type interval: int
:param interval: Check interval, in seconds
"""
self.model = compose.Pipeline(
preprocessing.MinMaxScaler(), anomaly.HalfSpaceTrees(
seed=42)) # tree.HoeffdingTreeClassifier(max_depth=10)
self.metric = metrics.ROCAUC() # metrics.Accuracy() #
self.my_id = my_id
self.t = request_topic
self.result_t = inference_topic
self.my_grp_id = group_id
self.result_t_p = 8
self.bootstrap_servers = bootstrap_servers
# self.list_of_partitions = list_of_partitions
self.tls = []
x = 0
for i in list_of_partitions:
self.tls.insert(x, TopicPartition(self.t, i))
x = x + 1
#self.tls=list_of_partitions
print(self.tls)
conf = {
'bootstrap.servers': bootstrap_servers,
'sasl.mechanism': 'PLAIN',
'security.protocol': 'SASL_SSL',
'ssl.ca.location': '/tmp/cacert.pem',
'sasl.username': '******',
'sasl.password':
'******',
# 'sasl.username': '******',
# 'sasl.password': '******',
# 'key.serializer': StringSerializer('utf_8'),
# 'value.serializer': StringSerializer('utf_8'),
'client.id': 'test-sw-1'
}
self.producer = Producer(conf)
conf = {
'bootstrap.servers': bootstrap_servers,
'sasl.mechanism': 'PLAIN',
'security.protocol': 'SASL_SSL',
'sasl.username': '******',
'sasl.password':
'******',
'ssl.ca.location': '/tmp/cacert.pem',
'group.id': group_id,
'auto.offset.reset': 'latest'
}
self.consumer = consumer = Consumer(conf)
self.consumer.assign(self.tls)
def __del__(self):
print('closing')
#self.consumer.close()
self.producer.flush()
self.producer.close()
def run(self):
### Set up model, metric, and starting timestamp for this model instance ###
model = self.model
metric = self.metric
print('MODEL and METRIC before any messages consumed:', model, metric)
start_consume_ts = time.time()
print('Start Consume Time ' + str(start_consume_ts))
######
try:
self.consumer.assign(self.tls)
i = 0
while (True):
msg = self.consumer.poll(timeout=1.0)
if msg is None: continue
if msg.error():
if msg.error().code() == KafkaError._PARTITION_EOF:
# End of partition event
sys.stderr.write(
'%% %s [%d] reached end at offset %d\n' %
(msg.topic(), msg.partition(), msg.offset()))
elif msg.error():
raise KafkaException(msg.error())
else:
message = loads(msg.value().decode("utf-8"))
ingest_ts = message[
'ingestTs'] # Ingestion time from Kafka
message.pop('ingestTs')
message_id = message['message_id'] # Message ID from Kafka
message.pop('message_id')
truth = float(
message['Class']) # True label supplied by Kafka
message.pop(
'Class'
) # Remove label remove passing the data for prediction (could combine with above)
message.pop('ignore')
x = message
for k, v in x.items():
x[k] = float(v)
y_hat = model.score_one(
x) #model.predict_one(x) # make a prediction
if (i % 1000 == 0):
print('x: ', x)
for k, v in x.items():
print(type(x[k]))
print('model: ', model)
print('y_hat: ', y_hat)
inference_ts = time.time()
metric = metric.update(truth, y_hat) # update the metric
model = model.learn_one(
x) # model.learn_one(x,y) # make the model learn
learn_ts = time.time()
out = {}
out['ingest_ts'] = ingest_ts
out['learn_ts'] = learn_ts
out['message_id'] = message_id
out['truth'] = truth
out['y_hat'] = y_hat
out['inference_ts'] = inference_ts
out['dur_evt_inf'] = inference_ts - ingest_ts
out['dur_start_inf'] = inference_ts - start_consume_ts
out['dur_inf_learn'] = learn_ts - inference_ts
out['model_metric'] = metric.get()
out['ignore'] = False
i = i + 1
# sprint(self.result_t)
k = str(i)
v = json.dumps(out).encode('utf-8')
self.producer.produce(self.result_t, value=v, key=k)
# self.producer.flush()
if (i % 1000 == 0):
print('sending to ' + self.result_t + ' ' +
str(self.result_t_p) + ' ' + str(out))
self.producer.flush()
if (i % 100000 == 0):
try:
pickle.dump(model, open("model.pkl", "wb"))
except:
os.remove("model.pkl")
pickle.dump(model, open("model.pkl", "wb"))
#self.producer.close()
finally:
# Close down consumer to commit final offsets.
self.consumer.close()
def main():
parser = argparse.ArgumentParser(
epilog="""Description:
Reidentification demo using any number of cameras:
Either camera can be used for registration or reidentification only, or for both.
Plays a video from a jpeg topic,
visualizes head detection with a gray bounding box around a head.
When a detection is identified, changes the bounding box color to orange
and writes the dwell time, age and ID (derived from the reid MS ID) above the heads.
Displays ('-d') or stores ('-o') the result of this demo in kafka topics.
Required topics (example):
- <prefix>.cam.0.original.Image.jpg
- <prefix>.cam.0.dets.ObjectDetectionRecord.json
- <prefix>.cam.0.frameinfo.FrameInfoRecord.json
- <prefix>.cam.0.ages.AgeRecord.json
- <prefix>.cam.1.original.Image.jpg
- <prefix>.cam.1.dets.ObjectDetectionRecord.json
- <prefix>.cam.1.frameinfo.FrameInfoRecord.json
- <prefix>.cam.1.ages.AgeRecord.json
...
- <prefix>.cam.1.reids.ReidRecord.json
""",
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("broker",
help="The name of the kafka broker.",
type=str)
parser.add_argument("prefix",
help="Prefix of topics (base|skeleton).",
type=str)
parser.add_argument('-d', "--display", action='store_true')
parser.add_argument('-o',
'--output',
help='write output image into kafka topic',
action='store_true')
parser.add_argument('text',
help='Text to display (age|dwell_time|both).',
type=str)
args = parser.parse_args()
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
overlay = cv2.imread('resources/powered_by_white.png',
cv2.IMREAD_UNCHANGED)
# Prepare the topics to read
input_topics = [
f"{args.prefix}.cam.{id}.{topic_postfix}" for id in CAMERA_TOPIC_IDS
for topic_postfix in TOPIC_POSTFIXES
]
reid_topics = [
f"{args.prefix}.cam.{id}.{topic_postfix}" for id in REID_TOPIC_IDS
for topic_postfix in REID_TOPIC_POSTFIXES
]
consumable_topics = list(map(TopicInfo, input_topics)) \
+ (list(map(lambda t: TopicInfo(t, drop=False), reid_topics)))
# TODO (when names via person stream): Remove this consumer
reg_consumer = Consumer({
'bootstrap.servers': args.broker,
'group.id': 'multicamreid_reg',
'auto.offset.reset': 'earliest'
})
reg_consumer.assign(
[TopicPartition(topic="named.records.json", partition=0, offset=0)])
output_topics = dict((id, f"{args.prefix}.cam.{id}.{OUTPUT_TOPIC_POSTFIX}")
for id in CAMERA_TOPIC_IDS)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(broker=args.broker,
groupid="detection",
topics_infos=consumable_topics,
latency_ms=200,
commit_interval_sec=None,
group_by_time=True)
registrations: Dict[str, Registration] = {}
i = 0
inner_id = 0
scaling = 1.0
for msgs in consumer.getMessages():
k = -1
for time, v in message_list_to_frame_structure(msgs).items():
message = v.get(args.prefix, {})
# Collect Reid records
reid_records = {}
for reid_id in REID_TOPIC_IDS:
reid_message = message.get(reid_id, {})
reid_records.update(reid_message.get("reid", {}))
# Process the image
for topic_key, topic_message in filter(
lambda t: t[0] not in REID_TOPIC_IDS, message.items()):
img = topic_message.get("image", {})
if not isinstance(img, np.ndarray):
continue
head_detections = topic_message.get("head_detection", {})
# Set the image scale
shape_orig = head_detections.pop("image", {})
if shape_orig:
scaling = img.shape[1] / shape_orig["frame_info"]["columns"]
# Processing the detections of the image
for detection_key, detection_record in head_detections.items():
object_detection_record = detection_record.get(
"bounding_box", {})
if not object_detection_record:
continue
key_to_display = ""
color = COLOR_DARK_GREY
face_detection = detection_record.get("unknown", {})
if face_detection:
color = COLOR_LIGHT_GREY
age = None
age_detection_record = detection_record.get("age", {})
if age_detection_record:
age = age_detection_record["age"]
if args.text == "age" or args.text == "both":
key_to_display = f"Age: {age}" if age else ""
# Reidentification received for the detection
reid_records_for_det = reid_records.get(detection_key, {})
if reid_records_for_det:
for reid_record in filter(lambda r: "reid_event" in r,
reid_records_for_det):
# We only use the first [0] identified face now
reid_key = reid_record["reid_event"]["match_list"][
0]["id"]["first_detection_key"]
registered = registrations.get(reid_key, None)
if registered:
age_to_display = ""
if age:
registered.addAge(age)
if args.text == "age" or args.text == "both":
age_to_display = f"; Age: {registered.age:d}" if age else ""
# Calculate the dwell time if required
dwell_time_display = ""
if args.text == "dwell_time" or args.text == "both":
detection_time = reid_record["reid_event"][
"match_list"][0]["id"][
"first_detection_time"]
dwell_time = time - int(detection_time)
dwell_time_display = f"; Dwell time: {dwell_time}ms"
color = COLOR_ORANGE
name_to_display = registered.name if registered.name else f"ID: {registered.id}"
key_to_display = f"{name_to_display}{age_to_display}{dwell_time_display}"
else:
inner_id += 1
registrations[reid_key] = Registration(
id=inner_id)
if age:
registrations[reid_key].addAge(age)
# Update the technical naming topic
# TODO (when names via person stream): remove
producer.produce(
"detected.records.json",
key=str(reid_key).encode("utf-8"),
value=(str(inner_id) +
";").encode("utf-8"),
timestamp=time)
# Read the technical naming topic
# TODO (when names via person stream): remove
reg_msg = reg_consumer.poll(0.01)
if reg_msg is not None:
try:
key = reg_msg.key().decode("utf-8")
name = reg_msg.value().decode("utf-8")
# Update the person name
reg_to_update = registrations.get(key)
if reg_to_update:
reg_to_update.addName(name)
else:
registrations[key] = Registration(name=name)
except:
print(
"Decoding entry of the named.records topic failed.",
flush=True)
# draw text above bounding box
img = draw_nice_text(
canvas=img,
text=key_to_display,
bounding_box=object_detection_record["bounding_box"],
color=color,
scale=scaling)
# draw bounding_box
img = draw_nice_bounding_box(
canvas=img,
bounding_box=object_detection_record["bounding_box"],
color=color,
scaling=scaling)
# draw ultinous logo
img = draw_overlay(canvas=img,
overlay=overlay,
position=Position.BOTTOM_RIGHT,
scale=scaling)
# produce output topic
if args.output:
out_topic = output_topics.get(topic_key)
producer.produce(out_topic,
value=encode_image_to_message(img),
timestamp=time)
producer.poll(0)
if i % 1000 == 0:
producer.flush()
i += 1
# display #
if args.display:
cv2.imshow(f"DEMO Camera {topic_key}", img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
def main():
parser = argparse.ArgumentParser(
epilog="""Description:
Plays a video from a jpeg topic, visualizes the head detections and tracks.
Displays the result on screen ('-d') or stores result in kafka ('-o').
Required topics:
- <prefix>.cam.0.original.Image.jpg
- <prefix>.cam.0.dets.ObjectDetectionRecord.json
- <prefix>.cam.0.tracks.TrackChangeRecord.json
""",
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("broker",
help="The name of the kafka broker.",
type=str)
parser.add_argument("prefix",
help="Prefix of topics (base|skeleton).",
type=str)
parser.add_argument('-f', "--full_screen", action='store_true')
parser.add_argument('-d', "--display", action='store_true')
parser.add_argument('-o',
'--output',
help='write output image into kafka topic',
action='store_true')
args = parser.parse_args()
if not args.display and not args.output:
parser.error(
"Missing argument: -d (display output) or -o (write output to kafka) is needed"
)
if args.output:
producer = Producer({'bootstrap.servers': args.broker})
overlay = cv2.imread('resources/powered_by_white.png',
cv2.IMREAD_UNCHANGED)
image_topic = f"{args.prefix}.cam.0.original.Image.jpg"
detection_topic = f"{args.prefix}.cam.0.dets.ObjectDetectionRecord.json"
track_topic = f"{args.prefix}.cam.0.tracks.TrackChangeRecord.json"
output_topic_name = f"{args.prefix}.cam.0.tracks.Image.jpg"
# handle full screen
window_name = "DEMO: Head detection"
if args.full_screen:
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
# read message, draw and display them
consumer = TimeOrderedGeneratorWithTimeout(args.broker, "detection", [
TopicInfo(image_topic),
TopicInfo(track_topic, drop=False),
TopicInfo(detection_topic)
], 100, None, True)
i = 0
tracks: DefaultDict[Any, Track] = defaultdict(lambda: Track(next(colors)))
for msgs in consumer.getMessages():
for time, v in message_list_to_frame_structure(msgs).items():
for track_key, track_val in v[args.prefix]["0"]["track"].items():
if track_val["end_of_track"]:
if track_key in tracks:
del tracks[track_key]
continue
point = track_val["point"]["x"], track_val["point"]["y"]
tracks[track_key].add_point(point)
img = v[args.prefix]["0"]["image"]
if type(img) == np.ndarray:
# draw bounding_box
for head_detection in v[args.prefix]["0"]["head_detection"]:
object_detection_record = v[args.prefix]["0"][
"head_detection"][head_detection]["bounding_box"]
if object_detection_record["type"] == "PERSON_HEAD":
img = draw_nice_bounding_box(
img, object_detection_record["bounding_box"],
(10, 95, 255))
for t_key, t in tracks.items():
cv2.polylines(img=img,
pts=[np.array(t.points, np.int32)],
isClosed=False,
color=t.color,
thickness=3)
# draw ultinous logo
img = draw_overlay(img, overlay, Position.BOTTOM_RIGHT)
# produce output topic
if args.output:
producer.produce(output_topic_name,
value=encode_image_to_message(img),
timestamp=time)
producer.poll(0)
if i % 100 == 0:
producer.flush()
i = 0
i += 1
# display
if args.display:
cv2.imshow(window_name, img)
k = cv2.waitKey(33)
if k == 113: # The 'q' key to stop
break
elif k == -1: # normally -1 returned,so don't print it
continue
else:
print(f"Press 'q' key for EXIT!")
class KafkaConn:
def __init__(self):
self.consume_lock = ConsumerState.PARTITIONS_UNASSIGNED
# callback for consumer partition assignment,
# removes lock for actual consumption
def callback_on_assignment(self, consumer, partitions):
self.consume_lock = ConsumerState.PARTITIONS_ASSIGNED
offsets = consumer.committed(partitions)
log.info(f"Assignment: {offsets}")
def check_connection(self):
def check_callback(error, event):
if error:
if error.code() == KafkaError._MSG_TIMED_OUT:
log.error(
"This Timout might indicate the broker is down or connection is misconfigured"
)
log.error(f"Error while producing initial msg: {error}")
raise KafkaSetupError()
config = ConfigFactory(kafka_client="producer").config
config["delivery.timeout.ms"] = "3000" # 3 seconds
prod = Producer(config)
prod.produce("pyrandall", "starting simulate", callback=check_callback)
prod.flush() # block until callback is called
def prod_reporter(self, error, event):
if error:
log.error(f"Error producing the event: {error}")
else:
log.info(f"Event produced, topic: {event.topic()}, \
partition: {event.partition()}")
def produce_message(self, topic, body, headers=None, partition_key=None):
if headers is None:
msg_headers = {}
self._produce(topic, body, partition_key, msg_headers)
self.producer.flush()
def init_producer(self):
log.info("starting produce")
kafka_config_producer = ConfigFactory(kafka_client="producer")
config = kafka_config_producer.config
log.info("kafka config for produce %s", config)
self.check_connection()
self.producer = Producer(config)
def _produce(self, topic, msg, partition_key=None, headers=None):
try:
if partition_key:
self.producer.produce(topic,
msg,
key=partition_key,
callback=self.prod_reporter)
else:
self.producer.produce(topic, msg, callback=self.prod_reporter)
print(".", end="")
except BufferError:
log.error(
"%% Local producer queue is full (%d messages \
awaiting delivery): try again",
len(self.producer),
)
# The consume function now contains a lock, the lock is removed when the
# partitions are assigned (max 60 seconds). After assignment the regular
# timeout are used. These should be set to a couple of seconds in the
# scenario itself .
def consume(self, topic, topic_timeout):
kafka_config_consumer = ConfigFactory(kafka_client="consumer")
config = kafka_config_consumer.config
log.info("kafka config for consume %s", config)
consumer = Consumer(config)
events = []
start_time = time.monotonic()
timeout_start_time = start_time
timeout_consumer = 10.0
# actual consumer starts now
# subscribe to 1 or more topics and define the callback function
# callback is only received after consumer.consume() is called!
consumer.subscribe([topic], on_assign=self.callback_on_assignment)
log.info(
f"Waiting for partition assignment ... (timeout at {timeout_consumer} seconds"
)
try:
while (time.monotonic() - timeout_start_time) < timeout_consumer:
# start consumption
messages = consumer.consume(timeout=0.1)
# check for partition assignment
if self.consume_lock == ConsumerState.PARTITIONS_UNASSIGNED:
# this should not happen but we are not 100% sure
if messages:
log.error("messages consumed but lock is unopened")
break
continue
# after partition assignment set the timeout again
# and reset the start time from which to determine timeout
# violation
elif self.consume_lock == ConsumerState.PARTITIONS_ASSIGNED:
timeout_start_time = time.monotonic()
timeout_consumer = topic_timeout
self.consume_lock = ConsumerState.TIMEOUT_SET
log.info("Lock has been opened, consuming ...")
# appened messages to the events list to be returned
if messages:
for msg in messages:
log.info(f"message at offset: {msg.offset()}, \
partition: {msg.partition()}, \
topic: {msg.topic()}")
# TODO: allow assertions to be on message headers etc.
# events.append({
# "key": msg.key,
# "headers": msg.headers,
# "value": msg.value()
# })
events.append(msg.value())
# only executed when while condition becomes false
else:
# at the end check if the partition assignment was achieved
if self.consume_lock != ConsumerState.TIMEOUT_SET:
log.error("No partition assignments received in time")
except KafkaException as e:
log.error(f"Kafka error: {e}")
pass
finally:
consumer.close()
end_time = time.monotonic()
log.debug(f"this cycle took: {(end_time - start_time)} seconds")
return events