不用 central broker 版本:
import pika, sys, os
import numpy as np
import cv2
import torch
from torch import nn
import torchvision
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
# load the COCO dataset category names
# we will use the same list for this notebook
COCO_INSTANCE_CATEGORY_NAMES = [
'__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()
class PersonDetector():
def __init__(self) -> None:
self.connection = pika.BlockingConnection(pika.ConnectionParameters(host='192.168.56.1', port=5672))
self.channel = self.connection.channel()
self.channel.exchange_declare(exchange="PersonDetector", exchange_type="topic", auto_delete=True, arguments={"output":["PersonDetector_output_text"]})
self.channel.queue_declare(queue='PersonDetector_input_image', exclusive=True)
self.channel.queue_bind(queue="PersonDetector_input_image", exchange="PersonDetector", routing_key=f"*.*.*.image")
# load model
print("Start loading model")
#model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
#model.eval()
print("Load model successfully")
def __callback(self, ch, method, properties, body):
if "PersonDetector" in method.routing_key:
pass
else:
routing_key_tokens = method.routing_key.split(".")
app_name = routing_key_tokens[0]
client_id = routing_key_tokens[1]
# preprocessing
img_bytes = np.frombuffer(body, dtype=np.uint8)
img = cv2.imdecode(img_bytes, 1)
cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img.astype(np.float32) / 255.0
img = torch.tensor(img)
img = img.permute(2, 0, 1)
# detect
pred = model([img])
pred_class = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].numpy())]
pred_boxes = [[(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]['boxes'].detach().numpy())]
pred_score = list(pred[0]['scores'].detach().numpy())
pred_t = [pred_score.index(x) for x in pred_score if x>0.7][-1]
pred_boxes = pred_boxes[:pred_t+1]
pred_class = pred_class[:pred_t+1]
for cls in pred_class:
if cls == "person":
print(f"person detect!!!")
def run(self):
self.channel.basic_consume(queue='PersonDetector_input_image', on_message_callback=self.__callback, auto_ack=True)
print(' [*] Waiting for messages. To exit press CTRL+C')
self.channel.start_consuming()
if __name__ == '__main__':
try:
detector = PersonDetector()
detector.run()
except KeyboardInterrupt:
print('Interrupted')
try:
sys.exit(0)
except SystemExit:
os._exit(0)
==========
import pika
import pandas as pd
import numpy as np
connection = pika.BlockingConnection(pika.ConnectionParameters(host='192.168.56.1', port=5672))
channel = connection.channel()
with open(f"C:\\Users\\Cherry\\Downloads\\image.jpg", "rb") as f:
data = f.read()
channel.basic_publish(exchange='PersonDetector',
routing_key=f'PersonDetection.client0.null.image',
body=data)
connection.close()
print("finish")
