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)