print("Distance detected by supersonic:", distance)
if distance != None and distance < 20:
robot.backward(0.35)
time.sleep(0.5)
robot.stop()
continue
# Capture frame-by-frame
frame = cap.read(
) # ret = 1 if the video is captured; frame is the image
# Our operations on the frame come here
img = cv2.flip(frame, 1) # flip left-right
img = cv2.flip(img, 0) # flip up-down
boxes, confs, clss = trt_yolo.detect(img, 0.8)
print("boxes:", boxes)
print("confs:", confs)
print("clss:", clss)
# Find target
target_clss_index = clss < 3
target_boxes = boxes[target_clss_index]
if len(target_boxes) > 0:
print("Found target ", len(target_boxes),
" and set barrier_center = -1")
barrier_center = -1
# Find barriers
barriers_clss_index = clss > 2
class YOLOv4:
def __init__(self,
frame_info_pool=None,
detection_result_pool=None,
final_result_pool=None,
model='yolov4-416',
dataset='obstacle'):
"""
:param model: model name
:param category_num:
"""
self.results = dict()
self.model_name = model
self.dataset = dataset
self.frame_info_pool = frame_info_pool
self.detection_result_pool = detection_result_pool
self.final_result_pool = final_result_pool
if dataset == "coco":
category_num = 80
else:
category_num = 15
self.cls_dict = get_cls_dict(category_num)
yolo_dim = model.split('-')[-1]
if 'x' in yolo_dim:
dim_split = yolo_dim.split('x')
if len(dim_split) != 2:
raise SystemExit('ERROR: bad yolo_dim (%s)!' % yolo_dim)
w, h = int(dim_split[0]), int(dim_split[1])
else:
h = w = int(yolo_dim)
if h % 32 != 0 or w % 32 != 0:
raise SystemExit('ERROR: bad yolo_dim (%s)!' % yolo_dim)
self.model = TrtYOLO(model, (h, w), category_num)
PrintLog.i("Object detection model is loaded - {}\t{}".format(
model, dataset))
def inference_by_image(self, image, confidence_threshold=0.1):
"""
:param image: input image
:param confidence_threshold: confidence/score threshold for object detection.
:return: bounding box(x1, y1, x2, y2), scores, classes
"""
boxes, scores, classes = self.model.detect(image, confidence_threshold)
results = []
for box, score, cls in zip(boxes, scores, classes):
results.append({
"label": [{
"description":
self.cls_dict.get(int(cls), 'CLS{}'.format(int(cls))),
"score":
float(score),
"class_idx":
int(cls)
}],
"position": {
"x": int(box[0]),
"y": int(box[1]),
"w": int(box[2] - box[0]),
"h": int(box[3] - box[1])
}
})
return results
def run(self):
PrintLog.i("Frame Number\tTimestamp")
while True:
if len(self.frame_info_pool) > 0:
frame_info = self.frame_info_pool.pop(0)
result = self.inference_by_image(frame_info["frame"])
detection_result = dict()
detection_result["cam_address"] = frame_info["cam_address"]
detection_result["timestamp"] = frame_info["timestamp"]
detection_result["frame_number"] = frame_info["frame_number"]
detection_result["results"] = []
detection_result["results"].append(
{"detection_result": result})
self.detection_result_pool.append(detection_result)
def main():
# initialize camera class
cam = Camera()
if not cam.is_opened():
raise SystemExit('EROR: failed to open camera!')
cls_dict = alprClassNames()
# Yolo dimensions (416x416)
yolo_dim = 416
h = w = int(yolo_dim)
# Initialized model and tools
cwd = os.getcwd()
model_yolo = str(cwd) + '/weights/yolov4-tiny-416.trt'
model_crnn = str(cwd) + '/weights/crnn.pth'
trt_yolo = TrtYOLO(model_yolo, (h,w), category_num=1) # category number is number of classes
crnn = alpr.AutoLPR(decoder='bestPath', normalise=True)
crnn.load(crnn_path=model_crnn)
open_window(WINDOW_NAME, TITLE, cam.img_width, cam.img_height)
vis = BBoxVisualization(cls_dict)
# Loop and detect
full_scrn = False
fps = 0.0
tic = time.time()
while True:
if cv2.getWindowProperty(WINDOW_NAME, 0) < 0:
break
img = cam.read()
if img is None:
break
# Detect car plate
boxes confs, clss = trt_yolo.detect(img, conf_th=0.5)
# Crop and preprocess car plate
cropped = vis.crop_plate(img, boxes, confs, clss)
# Recognize car plate
lp_plate = ''
fileLocate = str(cwd) + '/detections/detection1.jpg'
if os.path.exists(fileLocate):
lp_plate = lpr.predict(fileLocate)
# Draw boxes and fps
img = vis.draw_bboxes(img, boxes, confs, clss, lp=lp_plate)
img = show_fps(img, fps)
# Show image
cv2.imshow(WINDOW_NAME, img)
# Calculate fps
toc = time.time()
curr_fps = 1.0 / (toc - tic)
fps = curr_fps if fps == 0.0 else (fps*0.95 + curr_fps*0.05)
tic = toc
# Exit key
key = cv2.waitKey(1)
if key == 27: # ESC key: quit program
break
# Release capture and destroy all windows
cam.release()
cv2.destroyAllWindows()
class yolov4(object):
def __init__(self):
""" Constructor """
self.bridge = CvBridge()
self.init_params()
self.init_yolo()
self.cuda_ctx = cuda.Device(0).make_context()
self.trt_yolo = TrtYOLO((self.model_path + self.model),
(self.h, self.w), self.category_num)
def __del__(self):
""" Destructor """
self.cuda_ctx.pop()
del self.trt_yolo
del self.cuda_ctx
def clean_up(self):
""" Backup destructor: Release cuda memory """
if self.trt_yolo is not None:
self.cuda_ctx.pop()
del self.trt_yolo
del self.cuda_ctx
def init_params(self):
""" Initializes ros parameters """
rospack = rospkg.RosPack()
package_path = rospack.get_path("yolov4_trt_ros")
self.video_topic = rospy.get_param("/video_topic", "/video_source/raw")
self.model = rospy.get_param("/model", "yolov4")
self.model_path = rospy.get_param("/model_path",
package_path + "/yolo/")
self.input_shape = rospy.get_param("/input_shape", "416")
self.category_num = rospy.get_param("/category_number", 80)
self.conf_th = rospy.get_param("/confidence_threshold", 0.5)
self.show_img = rospy.get_param("/show_image", True)
self.image_sub = rospy.Subscriber(self.video_topic,
Image,
self.img_callback,
queue_size=1,
buff_size=1920 * 1080 * 3)
self.detection_pub = rospy.Publisher("detections",
Detector2DArray,
queue_size=1)
self.overlay_pub = rospy.Publisher("/result/overlay",
Image,
queue_size=1)
def init_yolo(self):
""" Initialises yolo parameters required for trt engine """
if self.model.find('-') == -1:
self.model = self.model + "-" + self.input_shape
yolo_dim = self.model.split('-')[-1]
if 'x' in yolo_dim:
dim_split = yolo_dim.split('x')
if len(dim_split) != 2:
raise SystemExit('ERROR: bad yolo_dim (%s)!' % yolo_dim)
self.w, self.h = int(dim_split[0]), int(dim_split[1])
else:
self.h = self.w = int(yolo_dim)
if self.h % 32 != 0 or self.w % 32 != 0:
raise SystemExit('ERROR: bad yolo_dim (%s)!' % yolo_dim)
cls_dict = get_cls_dict(self.category_num)
self.vis = BBoxVisualization(cls_dict)
def img_callback(self, ros_img):
"""Continuously capture images from camera and do object detection """
tic = time.time()
# converts from ros_img to cv_img for processing
try:
cv_img = self.bridge.imgmsg_to_cv2(ros_img,
desired_encoding="bgr8")
rospy.logdebug("ROS Image converted for processing")
except CvBridgeError as e:
rospy.loginfo("Failed to convert image %s", str(e))
if cv_img is not None:
boxes, confs, clss = self.trt_yolo.detect(cv_img, self.conf_th)
cv_img = self.vis.draw_bboxes(cv_img, boxes, confs, clss)
toc = time.time()
fps = 1.0 / (toc - tic)
self.publisher(boxes, confs, clss)
if self.show_img:
cv_img = show_fps(cv_img, fps)
cv2.imshow("YOLOv4 DETECTION RESULTS", cv_img)
cv2.waitKey(1)
# converts back to ros_img type for publishing
try:
overlay_img = self.bridge.cv2_to_imgmsg(cv_img,
encoding="passthrough")
rospy.logdebug("CV Image converted for publishing")
self.overlay_pub.publish(overlay_img)
except CvBridgeError as e:
rospy.loginfo("Failed to convert image %s", str(e))
def publisher(self, boxes, confs, clss):
""" Publishes to detector_msgs
Parameters:
boxes (List(List(int))) : Bounding boxes of all objects
confs (List(double)) : Probability scores of all objects
clss (List(int)) : Class ID of all classes
"""
detection2d = Detector2DArray()
detection = Detector2D()
for i in range(len(boxes)):
# boxes : xmin, ymin, xmax, ymax
for _ in boxes:
detection.header.stamp = rospy.Time.now()
detection.header.frame_id = "camera" # change accordingly
detection.results.id = clss[i]
detection.results.score = confs[i]
detection.bbox.center.x = (boxes[i][2] - boxes[i][0]) / 2
detection.bbox.center.y = (boxes[i][3] - boxes[i][1]) / 2
detection.bbox.center.theta = 0.0 # change if required
detection.bbox.size_x = abs(boxes[i][0] - boxes[i][2])
detection.bbox.size_y = abs(boxes[i][1] - boxes[i][3])
detection2d.detections.append(detection)
self.detection_pub.publish(detection2d)
tic = time.time()
def closedCapCallback():
global cap
print("no connection in the stream, reconnecting")
cap = cv2.VideoCapture('rtsp://localhost:8554/stream_input')
sleep(0.5)
while True:
if not cap.isOpened():
print('error while subprocess not running')
closedCapCallback()
else:
_,frame = cap.read()
if not frame is None:
boxes, confs, clss = trt_yolo.detect(frame, 0.3)
frame = vis.draw_bboxes(frame, boxes, confs, clss)
cv2.putText(frame, "FPS: " + str(round(fps, 2)), (10, 50), font, 3, (0, 0, 0), 3)
toc = time.time()
curr_fps = 1.0 / (toc - tic)
fps = curr_fps if fps == 0.0 else (fps * 0.95 + curr_fps * 0.05)
tic = toc
ret_toRTSP, frame_toRTSP = cv2.imencode('.png', frame)
process.stdin.write(frame_toRTSP.tobytes())
else:
print('error while subprocess is running')
closedCapCallback()
