def display(self,
pprint=False,
show=False,
save=False,
render=False,
save_dir=''):
colors = color_list()
for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
str = f'image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
if pred is not None:
for c in pred[:, -1].unique():
n = (pred[:, -1] == c).sum() # detections per class
str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
if show or save or render:
for *box, conf, cls in pred: # xyxy, confidence, class
label = f'{self.names[int(cls)]} {conf:.2f}'
plot_one_box(box,
img,
label=label,
color=colors[int(cls) % 10])
img = Image.fromarray(img.astype(np.uint8)) if isinstance(
img, np.ndarray) else img # from np
if pprint:
print(str.rstrip(', '))
if show:
img.show(self.files[i]) # show
if save:
f = Path(save_dir) / self.files[i]
img.save(f) # save
print(f"{'Saving' * (i == 0)} {f},",
end='' if i < self.n - 1 else ' done.\n')
if render:
self.imgs[i] = np.asarray(img)
def display(self, pprint=False, show=False, save=False, render=False):
colors = color_list()
for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
str = f'Image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
if pred is not None:
for c in pred[:, -1].unique():
n = (pred[:, -1] == c).sum() # detections per class
str += f'{n} {self.names[int(c)]}s, ' # add to string
if show or save or render:
img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np
for *box, conf, cls in pred: # xyxy, confidence, class
# str += '%s %.2f, ' % (names[int(cls)], conf) # label
ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10]) # plot
if pprint:
print(str)
if show:
img.show(f'Image {i}') # show
if save:
f = f'results{i}.jpg'
str += f"saved to '{f}'"
img.save(f) # save
if render:
self.imgs[i] = np.asarray(img)
Class has functions that can get 3d coordinates of detections and create Marker Spheres for visualizations
"""
from yolov5 import YOLOv5
from yolov5.utils.plots import color_list, plot_one_box
from yolov5.utils.general import xyxy2xywh
import random
import rospy
import numpy as np
from time import time
from image_geometry import PinholeCameraModel
from sensor_msgs.msg import CameraInfo
from visualization_msgs.msg import Marker, MarkerArray
model_path = "/home/beast/yolov5/weights/yolov5s.pt"
device = "cuda"
colors = color_list()
class PeopleDetection:
""" People Detection class with useful functions for getting coordinates of detections"""
def __init__(self):
self._net = YOLOv5(model_path, device) #jetson.inference.detectNet("ssd-mobilenet-v2")
self.img = None
self.width = None
self.height = None
self.need_cam_info = True
self.camera_model = PinholeCameraModel()
self.marker = Marker()
self.prev_time = 0
self.marker_pub = rospy.Publisher("visualization_markers", Marker, queue_size=10)
self.camera_info = rospy.Subscriber('/camera/depth/camera_info', CameraInfo, self.info_callback)