def __init__(self, id, device, config, log):
self.id = id
self.log = log
self.frame = []
self.device = device
self.config = config
self.cnt = 0
self.cams = []
self.img_size = config['img_size']
self.nms_max_overlap = config['nms_max_overlap']
self.body_min_w = config['body_min_w']
self.conf_thres = config['conf_thres']
self.nms_thres = config['nms_thres']
self.max_hum_w = int(self.img_size / 2)
self.detector = Darknet(self.config['model_def'],
img_size=self.config['img_size']).to(
self.device)
self.detector.load_darknet_weights(self.config['weights_path'])
self.encoder = gdet.create_box_encoder(
self.config['model_filename'],
batch_size=self.config['batch_size'])
self._stopevent = threading.Event()
print("created ok")
threading.Thread.__init__(self)
def load_model(self, use_onnx, half_precision, cfg_path, weights_path):
"""
Load model weights into device memory.
"""
model = Darknet(cfg_path, self.target_size)
model.load_state_dict(
torch.load(weights_path, map_location=self.device)['model'])
model.to(self.device).eval()
if use_onnx:
model.fuse()
img = torch.zeros((1, 3) + self.target_size)
torch.onnx.export(model,
img,
'weights/export.onnx',
verbose=False,
opset_version=10)
import onnx
model = onnx.load('weights/export.onnx')
onnx.checker.check_model(model)
if half_precision and self.device.type != 'cpu':
model.half()
torch.backends.cudnn.benchmark = True
return model
required=True,
help="Input cfg (.cfg) file path (required)")
args = parser.parse_args()
if not os.path.isfile(args.weights):
raise SystemExit("Invalid weights file")
if not os.path.isfile(args.cfg):
raise SystemExit("Invalid cfg file")
return args.weights, args.cfg
pt_file, cfg_file = parse_args()
wts_file = pt_file.split(".pt")[0] + ".wts"
device = select_device("cpu")
model = Darknet(cfg_file).to(device)
model.load_state_dict(torch.load(pt_file, map_location=device)["model"])
model.to(device).eval()
with open(wts_file, "w") as f:
wts_write = ""
conv_count = 0
for k, v in model.state_dict().items():
if not "num_batches_tracked" in k:
vr = v.reshape(-1).cpu().numpy()
wts_write += "{} {} ".format(k, len(vr))
for vv in vr:
wts_write += " "
wts_write += struct.pack(">f", float(vv)).hex()
wts_write += "\n"
conv_count += 1
class GpuDevice(threading.Thread):
def __init__(self, id, device, config, log):
self.id = id
self.log = log
self.frame = []
self.device = device
self.config = config
self.cnt = 0
self.cams = []
self.img_size = config['img_size']
self.nms_max_overlap = config['nms_max_overlap']
self.body_min_w = config['body_min_w']
self.conf_thres = config['conf_thres']
self.nms_thres = config['nms_thres']
self.max_hum_w = int(self.img_size / 2)
self.detector = Darknet(self.config['model_def'],
img_size=self.config['img_size']).to(
self.device)
self.detector.load_darknet_weights(self.config['weights_path'])
self.encoder = gdet.create_box_encoder(
self.config['model_filename'],
batch_size=self.config['batch_size'])
self._stopevent = threading.Event()
print("created ok")
threading.Thread.__init__(self)
#t = threading.Thread(target=self._reader)
#t.daemon = True
def startCam(self, camConfig):
print('start video: ', camConfig['id'])
#id, url, borders, skipFrames, max_cosine_distance, nn_budget
cam = videoCapture.VideoCapture(camConfig, self.config)
self.cams.append(cam)
print(self.cams[0].id)
if (len(self.cams) == 1):
self.start()
def stopCam(self, id):
print('stop video: ', id)
self.cams[0].exit()
time.sleep(0.1)
del self.cams[0]
if (len(self.cams) == 0):
time.sleep(0.5)
self.kill()
def removeCam(self, cam_id):
if cam_id in self.cams:
del self.cams[cam_id]
def kill(self):
self.log.debug("kill cameras " + str(len(self.cams)))
self._stopevent.set()
for cam in self.cams:
cam.exit()
print("kill ok")
self.killed = True
def run(self):
print("start camera")
cnt_people_in = []
tr = True
while not self._stopevent.isSet():
start = time.time()
self.cnt += 1
frames = []
for cam in self.cams:
if self.cams[0]._stopevent.isSet():
#delete cam from cams list
pass
else:
frames.append(cam.read())
if (tr): cv2.imwrite("video/39_2.jpg", frames[0])
frame = frames[0]
#frame_cuda = transforms.ToTensor()(frame).unsqueeze(0)
frame_cuda = transforms.ToTensor()(frame).to(
self.device).unsqueeze(0)
with torch.no_grad():
obj_detec = self.detector(frame_cuda)
obj_detec = non_max_suppression(obj_detec, self.conf_thres,
self.nms_thres)
boxs = []
confs = []
for item in obj_detec:
if item is not None:
i = 0
for x1, y1, x2, y2, conf, cls_conf, cls_pred in item: #classes[int(cls_pred)]
wb = y2 - y1
if ((cls_pred == 0) and (wb < self.max_hum_w)
and (wb > self.body_min_w)):
boxs.append(np.array([y1, x1, y2, x2]))
# print("len=", len(boxs))
if (len(boxs)):
t_start2 = time.time()
features = self.encoder(frame, boxs)
detections = [
Detection(bbox, 1.0, feature)
for bbox, feature in zip(boxs, features)
]
boxes = np.array([d.tlwh for d in detections
]) # w and h replace by x2 y2
scores = np.array([d.confidence for d in detections])
indices = preprocessing.non_max_suppression(
boxes, self.nms_max_overlap, scores)
detections = [detections[i] for i in indices]
for i, cam in enumerate(self.cams):
cam.track(detections, frames[i])
img_size=args.img_size,
multiscale=False)
test_loader = DataLoader(test_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.n_cpu,
collate_fn=test_set.collate_fn)
else:
test_set = ImageFolder(args.image_folder, img_size=args.img_size)
test_loader = DataLoader(test_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.n_cpu)
print("3. Creating Model!")
model = Darknet(args.model_def).to(device=device)
if args.weights_path:
if args.weights_path.endswith(".pth"):
model.load_state_dict(torch.load(args.weights_path))
else:
model.load_darknet_weights(args.weights_path)
print("4. Create csvfile to save test results. ")
with open(args.save_path / args.log, 'w') as csvfile:
csv_writer = csv.writer(csvfile, delimiter='\t')
csv_writer.writerow('test results!')
print("5. Start Detection!")
model.eval()
def __init__(self):
# Load weights parameter
weights_name = rospy.get_param('~weights_name', 'yolov3.weights')
self.weights_path = os.path.join(package_path, 'models', weights_name)
rospy.loginfo("Found weights, loading %s", self.weights_path)
# Raise error if it cannot find the model
if not os.path.isfile(self.weights_path):
raise IOError(('{:s} not found.').format(self.weights_path))
# Load image parameter and confidence threshold
self.image_topic = rospy.get_param('~image_topic',
'/camera/rgb/image_raw')
self.confidence_th = rospy.get_param('~confidence', 0.7)
self.nms_th = rospy.get_param('~nms_th', 0.3)
# Load publisher topics
self.detected_objects_topic = rospy.get_param(
'~detected_objects_topic')
self.published_image_topic = rospy.get_param('~detections_image_topic')
# Load other parameters
config_name = rospy.get_param('~config_name', 'yolov3.cfg')
self.config_path = os.path.join(package_path, 'config', config_name)
classes_name = rospy.get_param('~classes_name', 'coco.names')
self.classes_path = os.path.join(package_path, 'classes', classes_name)
self.gpu_id = rospy.get_param('~gpu_id', 0)
self.network_img_size = rospy.get_param('~img_size', 416)
self.publish_image = rospy.get_param('~publish_image')
# Initialize width and height
self.h = 0
self.w = 0
rospy.loginfo("config path: " + self.config_path)
self.model = Darknet(self.config_path, img_size=self.network_img_size)
# Load net
self.model.load_weights(self.weights_path)
if torch.cuda.is_available():
rospy.loginfo("CUDA available, use GPU")
self.model.cuda()
else:
rospy.loginfo("CUDA not available, use CPU")
# if CUDA not available, use CPU
# self.checkpoint = torch.load(self.weights_path, map_location=torch.device('cpu'))
# self.model.load_state_dict(self.checkpoint)
self.model.eval() # Set in evaluation mode
rospy.loginfo("Deep neural network loaded")
# Load CvBridge
self.bridge = CvBridge()
# Load classes
self.classes = load_classes(
self.classes_path) # Extracts class labels from file
self.classes_colors = {}
# Define subscribers
self.image_sub = rospy.Subscriber(self.image_topic,
Image,
self.imageCb,
queue_size=1,
buff_size=2**24)
# Define publishers
self.pub_ = rospy.Publisher(self.detected_objects_topic,
BoundingBoxes,
queue_size=10)
self.pub_viz_ = rospy.Publisher(self.published_image_topic,
Image,
queue_size=10)
rospy.loginfo("Launched node for object detection")
class DetectorManager():
def __init__(self):
# Load weights parameter
weights_name = rospy.get_param('~weights_name', 'yolov3.weights')
self.weights_path = os.path.join(package_path, 'models', weights_name)
rospy.loginfo("Found weights, loading %s", self.weights_path)
# Raise error if it cannot find the model
if not os.path.isfile(self.weights_path):
raise IOError(('{:s} not found.').format(self.weights_path))
# Load image parameter and confidence threshold
self.image_topic = rospy.get_param('~image_topic',
'/camera/rgb/image_raw')
self.confidence_th = rospy.get_param('~confidence', 0.7)
self.nms_th = rospy.get_param('~nms_th', 0.3)
# Load publisher topics
self.detected_objects_topic = rospy.get_param(
'~detected_objects_topic')
self.published_image_topic = rospy.get_param('~detections_image_topic')
# Load other parameters
config_name = rospy.get_param('~config_name', 'yolov3.cfg')
self.config_path = os.path.join(package_path, 'config', config_name)
classes_name = rospy.get_param('~classes_name', 'coco.names')
self.classes_path = os.path.join(package_path, 'classes', classes_name)
self.gpu_id = rospy.get_param('~gpu_id', 0)
self.network_img_size = rospy.get_param('~img_size', 416)
self.publish_image = rospy.get_param('~publish_image')
# Initialize width and height
self.h = 0
self.w = 0
rospy.loginfo("config path: " + self.config_path)
self.model = Darknet(self.config_path, img_size=self.network_img_size)
# Load net
self.model.load_weights(self.weights_path)
if torch.cuda.is_available():
rospy.loginfo("CUDA available, use GPU")
self.model.cuda()
else:
rospy.loginfo("CUDA not available, use CPU")
# if CUDA not available, use CPU
# self.checkpoint = torch.load(self.weights_path, map_location=torch.device('cpu'))
# self.model.load_state_dict(self.checkpoint)
self.model.eval() # Set in evaluation mode
rospy.loginfo("Deep neural network loaded")
# Load CvBridge
self.bridge = CvBridge()
# Load classes
self.classes = load_classes(
self.classes_path) # Extracts class labels from file
self.classes_colors = {}
# Define subscribers
self.image_sub = rospy.Subscriber(self.image_topic,
Image,
self.imageCb,
queue_size=1,
buff_size=2**24)
# Define publishers
self.pub_ = rospy.Publisher(self.detected_objects_topic,
BoundingBoxes,
queue_size=10)
self.pub_viz_ = rospy.Publisher(self.published_image_topic,
Image,
queue_size=10)
rospy.loginfo("Launched node for object detection")
def imageCb(self, data):
# Convert the image to OpenCV
try:
self.cv_image = self.bridge.imgmsg_to_cv2(data, "rgb8")
except CvBridgeError as e:
print(e)
# Initialize detection results
detection_results = BoundingBoxes()
detection_results.header = data.header
detection_results.image_header = data.header
# Configure input
input_img = self.imagePreProcessing(self.cv_image)
# set image type
if (torch.cuda.is_available()):
input_img = Variable(input_img.type(torch.cuda.FloatTensor))
else:
input_img = Variable(input_img.type(torch.FloatTensor))
# Get detections from network
with torch.no_grad():
detections = self.model(input_img)
detections = non_max_suppression(detections, 80,
self.confidence_th, self.nms_th)
# Parse detections
if detections[0] is not None:
for detection in detections[0]:
# Get xmin, ymin, xmax, ymax, confidence and class
xmin, ymin, xmax, ymax, _, conf, det_class = detection
pad_x = max(self.h - self.w,
0) * (self.network_img_size / max(self.h, self.w))
pad_y = max(self.w - self.h,
0) * (self.network_img_size / max(self.h, self.w))
unpad_h = self.network_img_size - pad_y
unpad_w = self.network_img_size - pad_x
xmin_unpad = ((xmin - pad_x // 2) / unpad_w) * self.w
xmax_unpad = ((xmax - xmin) / unpad_w) * self.w + xmin_unpad
ymin_unpad = ((ymin - pad_y // 2) / unpad_h) * self.h
ymax_unpad = ((ymax - ymin) / unpad_h) * self.h + ymin_unpad
# Populate darknet message
detection_msg = BoundingBox()
detection_msg.xmin = xmin_unpad
detection_msg.xmax = xmax_unpad
detection_msg.ymin = ymin_unpad
detection_msg.ymax = ymax_unpad
detection_msg.probability = conf
detection_msg.Class = self.classes[int(det_class)]
# Append in overall detection message
detection_results.bounding_boxes.append(detection_msg)
# Publish detection results
self.pub_.publish(detection_results)
# Visualize detection results
if (self.publish_image):
self.visualizeAndPublish(detection_results, self.cv_image)
else:
rospy.loginfo("No detections available, next image")
return True
def imagePreProcessing(self, img):
# Extract image and shape
img = np.ascontiguousarray(img)
img = img.astype(float)
height, width, channels = img.shape
if (height != self.h) or (width != self.w):
self.h = height
self.w = width
# Determine image to be used
self.padded_image = np.zeros(
(max(self.h, self.w), max(self.h,
self.w), channels)).astype(float)
# Add padding
if (self.w > self.h):
self.padded_image[(self.w - self.h) // 2:self.h +
(self.w - self.h) // 2, :, :] = img
else:
self.padded_image[:, (self.h - self.w) // 2:self.w +
(self.h - self.w) // 2, :] = img
# Resize and normalize
input_img = resize(
self.padded_image,
(self.network_img_size, self.network_img_size, 3)) / 255.
# Channels-first
input_img = np.transpose(input_img, (2, 0, 1))
# As pytorch tensor
input_img = torch.from_numpy(input_img).float()
input_img = input_img[None]
return input_img
def visualizeAndPublish(self, output, imgIn):
# Copy image and visualize
imgOut = np.ascontiguousarray(imgIn)
font = cv2.FONT_HERSHEY_SIMPLEX
fontScale = 0.6
thickness = int(2)
for index in range(len(output.bounding_boxes)):
label = output.bounding_boxes[index].Class
x_p1 = output.bounding_boxes[index].xmin
y_p1 = output.bounding_boxes[index].ymin
x_p3 = output.bounding_boxes[index].xmax
y_p3 = output.bounding_boxes[index].ymax
confidence = output.bounding_boxes[index].probability
# Find class color
if label in self.classes_colors.keys():
color = self.classes_colors[label]
else:
# Generate a new color if first time seen this label
color = np.random.randint(0, 255, 3)
self.classes_colors[label] = color
# Create rectangle
start_point = (int(x_p1), int(y_p1))
end_point = (int(x_p3), int(y_p3))
lineColor = (int(color[0]), int(color[1]), int(color[2]))
cv2.rectangle(imgOut, start_point, end_point, lineColor, thickness)
text = ('{:s}: {:.3f}').format(label, confidence)
cv2.putText(imgOut, text, (int(x_p1), int(y_p1 + 20)), font,
fontScale, (255, 255, 255), thickness, cv2.LINE_AA)
# Publish visualization image
image_msg = self.bridge.cv2_to_imgmsg(imgOut, "rgb8")
self.pub_viz_.publish(image_msg)
train_loader = DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.n_cpu,
pin_memory=True,
collate_fn=train_set.collate_fn)
valid_loader = DataLoader(valid_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.n_cpu,
pin_memory=True,
collate_fn=valid_set.collate_fn)
print("3. Creating Model")
# Initiate model
model = Darknet(args.model_def).to(device)
model.apply(weights_init_normal)
# If specified we start from checkpoint
if args.pretrained_weights.endswith(".pth"):
model.load_state_dict(torch.load(args.pretrained_weights))
else:
model.load_darknet_weights(args.pretrained_weights)
print("4. Setting Optimization Solver")
optimizer = torch.optim.Adam(model.parameters(), lr=0.0002, weight_decay=1e-4)
exp_lr_scheduler_R = lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.2)
print("5. Start Tensorboard")
mag = math.sqrt(vX0 * vX0 + vY0 * vY0)
vX = vX0 / mag
vY = vY0 / mag
temp = vX
vX = -vY
vY = temp
z0 = (int(a[0] + vX0 / 2), int(a[1] + vY0 / 2))
z1 = (int(a[0] + vX0 / 2 - vX * length), int(a[1] + vY0 / 2 - vY * length))
cv2.line(frame, a, b, (255, 255, 0), 2)
cv2.arrowedLine(frame, z0, z1, (0, 255, 0), 2)
cv2.putText(frame, "Out", z1, 0, 1, (0, 255, 0), 1)
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
detector = Darknet(model_def, img_size=img_size).to(device)
detector.load_darknet_weights(weights_path)
tr_img = True
nn_budget = None
path_track = 20 # how many frames in path are saves
cnt_people_in = {}
#cnt_people_out = {}
skip_counter = 10
counter = 0
# human ids model
model_filename = 'models/mars-small128.pb'
# wget https://github.com/Qidian213/deep_sort_yolov3/raw/master/model_data/mars-small128.pb
encoder = gdet.create_box_encoder(model_filename, batch_size=64)