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
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)
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
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")
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)