def __init__(self, dataloder, path=".", batchSize=1, queueSize=1024, inference_steps=1):
# initialize the file video stream along with the boolean
# used to indicate if the thread should be stopped or not
self.det_model = Darknet(f"{path}/yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights(f"{path}/models/yolo/yolov3-spp.weights")
self.det_model.net_info["height"] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info["height"])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda()
self.det_model.eval()
# Only perform inference every this amount of frames
self.inference_steps = inference_steps
self.stopped = False
self.dataloder = dataloder
self.batchSize = batchSize
self.datalen = self.dataloder.length()
leftover = 0
if (self.datalen) % batchSize:
leftover = 1
self.num_batches = self.datalen // batchSize + leftover
# initialize the queue used to store frames read from
# the video file
self.Q = Queue(maxsize=queueSize)
def __init__(self, dataloder, obj_id, batchSize=1, queueSize=1024):
# initialize the file video stream along with the boolean
# used to indicate if the thread should be stopped or not
cfg_path = "yolo/cfg/yolov3-single.cfg"
weights_path = 'models/yolo/{:02d}.weights'.format(obj_id)
self.det_model = Darknet(cfg_path, reso=int(opt.inp_dim))
self.det_model.load_weights(weights_path)
print("Loading YOLO cfg from", cfg_path)
print("Loading YOLO weights from", weights_path)
self.det_model.eval()
self.det_model.net_info['height'] = opt.inp_dim #input_dimension
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda()
self.det_model.eval()
self.stopped = False
self.dataloder = dataloder
self.batchSize = batchSize
self.datalen = self.dataloder.length()
leftover = 0
if (self.datalen) % batchSize:
leftover = 1
self.num_batches = self.datalen // batchSize + leftover
# initialize the queue used to store frames read from
# the video file
if opt.sp:
self.Q = Queue(maxsize=queueSize)
else:
self.Q = mp.Queue(maxsize=queueSize)
def __init__(self,
confidence=0.5,
nms_thresh=0.4,
resolution=416,
weights_path='yolo/weights/yolov3.weights',
cfg_path='yolo/cfg/yolov3.cfg',
num_classes=80,
names_path='yolo/data/coco.names'):
self.confidence = confidence
self.nms_thesh = nms_thresh
self.weightsfile = weights_path
self.cfgfile = cfg_path
self.CUDA = torch.cuda.is_available()
self.num_classes = num_classes
self.classes = load_classes(names_path)
self.model = Darknet(self.cfgfile)
self.model.load_weights(self.weightsfile)
self.model.net_info["height"] = resolution
self.inp_dim = int(self.model.net_info["height"])
#Check if resolution is multiple of 32
assert self.inp_dim % 32 == 0
assert self.inp_dim > 32
# If there's a GPU availible, put the model on GPU
if self.CUDA:
self.model.cuda()
# Set model in evaluation mode
self.model.eval()
def __init__(self, cuda_id=0, fast_yolo=False):
self.time_det = 0.0
self.time_run = 0.0
self.cuda_id = cuda_id
self.target_kps = [5, 6, 7, 8, 9, 10]
# Load yolo detection model
print('Loading YOLO model..')
if fast_yolo:
self.det_model = Darknet('./AlphaPose/yolo/cfg/yolov3-tiny.cfg', self.cuda_id)
self.det_model.load_weights('./AlphaPose/models/yolo/yolov3-tiny.weights')
else:
self.det_model = Darknet('./AlphaPose/yolo/cfg/yolov3.cfg', self.cuda_id)
self.det_model.load_weights('./AlphaPose/models/yolo/yolov3.weights')
self.det_model.cuda(self.cuda_id)
self.det_model.eval()
# Load pose model
print('Loading Alphapose pose model..')
pose_dataset = Mscoco()
if args.fast_inference:
self.pose_model = InferenNet_fast(4 * 1 + 1, pose_dataset)
else:
self.pose_model = InferenNet(4 * 1 + 1, pose_dataset)
self.pose_model.cuda(self.cuda_id)
self.pose_model.eval()
def __init__(self, path, batchSize=4, queueSize=256):
# initialize the file video stream along with the boolean
# used to indicate if the thread should be stopped or not
self.det_model = Darknet("yolo/cfg/yolov3-spp.cfg",
reso=int(opt.inp_dim))
self.det_model.load_weights('models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda()
self.det_model.eval()
self.stream = cv2.VideoCapture(path)
assert self.stream.isOpened(), 'Cannot capture source'
self.stopped = False
self.batchSize = batchSize
self.datalen = int(self.stream.get(cv2.CAP_PROP_FRAME_COUNT))
leftover = 0
if (self.datalen) % batchSize:
leftover = 1
self.num_batches = self.datalen // batchSize + leftover
# initialize the queue used to store frames read from
# the video file
self.Q = Queue(maxsize=queueSize)
def __init__(self):
self.confidence = 0.7
self.nms_thesh = 0.4
self.resolution = 640
self.scales = "1,2,3"
self.confidence = float(self.confidence)
self.nms_thesh = float(self.nms_thesh)
self.CUDA = torch.cuda.is_available()
self.num_classes = 80
self.classes = load_classes('yolo/data/coco.names')
print("Loading network.....")
self.model_detect = Darknet('cfg/yolov3.cfg')
self.model_detect.load_weights('yolo/yolov3.weights')
print("Network successfully loaded")
self.model_detect.net_info["height"] = self.resolution
self.inp_dim = int(self.model_detect.net_info["height"])
assert self.inp_dim % 32 == 0
assert self.inp_dim > 32
if self.CUDA:
self.model_detect.cuda()
self.model_detect.eval()
self.colors = pkl.load(open("yolo/pallete", "rb"))
def __init__(self):
super().__init__()
self.cfgfile = "yolo/cfg/yolov3.cfg"
self.weightsfile = "yolo/yolov3.weights"
self.num_classes = 80
self.confidence = 0.25
self.nms_thesh = 0.4
self.reso = 160
self.start = 0
self.CUDA = torch.cuda.is_available()
self.num_classes = 80
self.bbox_attrs = 5 + self.num_classes
self.model = Darknet(self.cfgfile)
self.model.load_weights(self.weightsfile)
self.model.net_info["height"] = self.reso
self.inp_dim = int(self.model.net_info["height"])
assert self.inp_dim % 32 == 0
assert self.inp_dim > 32
if self.CUDA:
self.model.cuda()
# Switch to “evaluate” mode before predictions
self.model.eval()
self.classes = load_classes('yolo/data/coco.names')
self.colors = pkl.load(open("yolo/pallete", "rb"))
def __init__(self, dataloder, batchSize=1, queueSize=1024, use_boxGT=False, gt_json=''):
# initialize the file video stream along with the boolean
# used to indicate if the thread should be stopped or not
self.det_model = Darknet(opt.yolo_model_cfg)
self.det_model.load_weights(opt.yolo_model_path)
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda()
self.det_model.eval()
self.stopped = False
self.dataloder = dataloder
self.batchSize = batchSize
self.datalen = self.dataloder.length()
leftover = 0
if (self.datalen) % batchSize:
leftover = 1
self.num_batches = self.datalen // batchSize + leftover
# initialize the queue used to store frames read from
# the video file
if opt.sp:
self.Q = Queue(maxsize=queueSize)
else:
self.Q = mp.Queue(maxsize=queueSize)
self.use_boxGT = use_boxGT
if use_boxGT:
print('loading grondtruth box.')
self.box_gt = box_gt_class(gt_json)
def __init__(self, dataloder, batchSize=1, queueSize=1024):
# initialize the file video stream along with the boolean
# used to indicate if the thread should be stopped or not
self.det_model = Darknet("yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights('models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda()
self.det_model.eval()
self.stopped = False
self.dataloder = dataloder
self.batchSize = batchSize
self.datalen = self.dataloder.length()
leftover = 0
if (self.datalen) % batchSize:
leftover = 1
self.num_batches = self.datalen // batchSize + leftover
# initialize the queue used to store frames read from
# the video file
if opt.sp:
self.Q = Queue(maxsize=queueSize)
else:
self.Q = mp.Queue(maxsize=queueSize)
def __init__(self):
self.det_model = Darknet("yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights('models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda()
self.det_model.eval()
class ObjectDetection(object):
def __init__(self, batchSize=1):
self.det_model = Darknet("yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights('models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.eval()
self.stopped = False
self.batchSize = batchSize
def process(self, img, orig_img, im_name, im_dim_list):
with torch.no_grad():
# Human Detection
img = img
prediction = self.det_model(img, CUDA=False)
# NMS process
dets = dynamic_write_results(prediction,
opt.confidence,
opt.num_classes,
nms=True,
nms_conf=opt.nms_thesh)
if isinstance(dets, int) or dets.shape[0] == 0:
return orig_img[0], im_name[0], None, None, None, None, None
dets = dets.cpu()
im_dim_list = torch.index_select(im_dim_list, 0, dets[:, 0].long())
scaling_factor = torch.min(self.det_inp_dim / im_dim_list,
1)[0].view(-1, 1)
# coordinate transfer
dets[:, [1, 3]] -= (self.det_inp_dim - scaling_factor *
im_dim_list[:, 0].view(-1, 1)) / 2
dets[:, [2, 4]] -= (self.det_inp_dim - scaling_factor *
im_dim_list[:, 1].view(-1, 1)) / 2
dets[:, 1:5] /= scaling_factor
for j in range(dets.shape[0]):
dets[j, [1, 3]] = torch.clamp(dets[j, [1, 3]], 0.0,
im_dim_list[j, 0])
dets[j, [2, 4]] = torch.clamp(dets[j, [2, 4]], 0.0,
im_dim_list[j, 1])
boxes = dets[:, 1:5]
scores = dets[:, 5:6]
boxes_k = boxes[dets[:, 0] == 0]
if isinstance(boxes_k, int) or boxes_k.shape[0] == 0:
return orig_img[0], im_name[0], None, None, None, None, None
inps = torch.zeros(boxes_k.size(0), 3, opt.inputResH, opt.inputResW)
pt1 = torch.zeros(boxes_k.size(0), 2)
pt2 = torch.zeros(boxes_k.size(0), 2)
return orig_img[0], im_name[0], boxes_k, scores[dets[:, 0] ==
0], inps, pt1, pt2
def gen_torch_weight(infname, outfname):
"""Convert darknet style model file to PyTorch format.
Paramaters
----------
infname : Darknet style model file(*.weights).
outfname : PyTorch style model file(*.pkl).
"""
net = Darknet('cfg/yolov3.cfg')
net.load_darknet_weights(infname)
torch.save(net.state_dict(), outfname)
def init_model(args):
scales = args.scales
images = args.images
batch_size = int(args.bs)
confidence = float(args.confidence)
nms_thesh = float(args.nms_thresh)
start = 0
num_classes = 80
classes = load_classes('yolo/data/coco.names')
print("classes")
print(classes)
# Set up the neural network
print("Loading network.....")
model = Darknet(args.cfgfile)
model.load_weights(args.weightsfile)
print("Network successfully loaded")
model.net_info["height"] = args.reso
inp_dim = int(model.net_info["height"])
assert inp_dim % 32 == 0
assert inp_dim > 32
# If there's a GPU availible, put the model on GPU
if CUDA:
model.cuda()
# Set the model in evaluation mode
model.eval()
return model
def __init__(self):
# initialize the file video stream along with the boolean
# used to indicate if the thread should be stopped or not
self.det_model = Darknet("./yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights('./models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda(torchCuda)
self.det_model.eval()
def __init__(self, batchSize=1):
self.det_model = Darknet("yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights('models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.eval()
self.stopped = False
self.batchSize = batchSize
def __init__(self, cuda_id=0, fast_yolo=False):
self.time_det = 0.0
self.time_run = 0.0
self.num_joints = 17
self.target_kps = [5, 6, 7, 8, 9, 10]
# Load yolo detection model
print('Loading YOLO model..')
if fast_yolo:
self.det_model = Darknet('./AlphaPose/yolo/cfg/yolov3-tiny.cfg')
self.det_model.load_weights(
'./AlphaPose/models/yolo/yolov3-tiny.weights')
else:
self.det_model = Darknet("./AlphaPose/yolo/cfg/yolov3.cfg")
self.det_model.load_weights(
'./AlphaPose/models/yolo/yolov3.weights')
self.det_model.cuda()
self.det_model.eval()
cfg_file = 'MSRAPose/experiments/coco/resnet50/256x192_d256x3_adam_lr1e-3.yaml'
model_file = 'MSRAPose/models/pytorch/pose_coco/pose_resnet_50_256x192.pth.tar'
# update config
update_config(cfg_file)
config.TEST.MODEL_FILE = model_file
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
# load pre-trained model
self.model = eval('models_msra.' + config.MODEL.NAME +
'.get_pose_net')(config, is_train=False)
print('Loading MSRA pose model..')
print('=> loading model from {}'.format(config.TEST.MODEL_FILE))
self.model.load_state_dict(torch.load(config.TEST.MODEL_FILE))
gpus = [int(i) for i in config.GPUS.split(',')]
self.model = torch.nn.DataParallel(self.model, device_ids=gpus).cuda()
self.model.eval()
# image transform
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
def __init__(self, dataloder, batchSize=1):
# initialize the file video stream along with the boolean
# used to indicate if the thread should be stopped or not
self.det_model = Darknet("yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights('models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda()
self.det_model.eval()
self.dataloader = dataloder
self.stopped = False
self.batchSize = batchSize
def load_yolo_model(args):
print('loading yolo model ...')
det_model = Darknet("yolo/cfg/yolov3-spp.cfg")
det_model.load_weights('models/yolo/yolov3-spp.weights')
det_model.net_info['height'] = args.inp_dim
det_inp_dim = int(det_model.net_info['height'])
assert det_inp_dim % 32 == 0
assert det_inp_dim > 32
det_model.cuda()
det_model.eval()
return det_model, det_inp_dim
def load_model(self):
args = self.detector_opt
print('Loading YOLO model..')
self.model = Darknet(self.model_cfg)
self.model.load_weights(self.model_weights)
self.model.net_info['height'] = self.inp_dim
if args:
if len(args.gpus) > 1:
self.model = torch.nn.DataParallel(self.model, device_ids=args.gpus).to(args.device)
else:
self.model.to(args.device)
else:
self.model.cuda()
self.model.eval()
def load_model(opt):
pose_dataset = Mscoco()
pose_model = InferenNet_fast(4 * 1 + 1, pose_dataset)
det_model = Darknet("yolo/cfg/yolov3-spp.cfg")
det_model.load_weights('models/yolo/yolov3-spp.weights')
det_model.net_info['height'] = opt.inp_dim
pose_model.cuda()
pose_model.eval()
det_model.cuda()
det_model.eval()
return det_model, pose_model
def __init__(self, batchSize=1, queueSize=1, size=100, device=0):
## camera stream
self.stream = cv2.VideoCapture(device)
assert self.stream.isOpened(), 'Cannot capture from camera'
self.stream.set(cv2.CAP_PROP_BUFFERSIZE, 1)
self.inp_dim = int(opt.inp_dim)
## yolo model
self.det_model = Darknet("joints_detectors/Alphapose/yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights('joints_detectors/Alphapose/models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda()
self.det_model.eval()
self.batchSize = batchSize
self.datalen = 1
leftover = 0
if (self.datalen) % batchSize:
leftover = 1
self.num_batches = self.datalen // batchSize + leftover
## alphapose model
fast_inference = True
pose_dataset = Mscoco()
if fast_inference:
self.pose_model = InferenNet_fast(4 * 1 + 1, pose_dataset)
else:
self.pose_model = InferenNet(4 * 1 + 1, pose_dataset)
self.pose_model.cuda()
self.pose_model.eval()
## 2d plotting
self.fig_in = plt.figure(figsize=(size , size))
self.ax_in = self.fig_in.add_subplot(1, 1, 1)
self.ax_in.get_xaxis().set_visible(False)
self.ax_in.get_yaxis().set_visible(False)
self.ax_in.set_axis_off()
self.ax_in.set_title('Input')
self.initialized = False
self.size=size
thismanager = get_current_fig_manager()
thismanager.window.wm_geometry("+0-1000")
def __init__(self, dataloder, batchSize=1, queueSize=1024):
# initialize the file video stream along with the boolean
# used to indicate if the thread should be stopped or not
self.det_model = Darknet("yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights('models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cpu()
self.det_model.eval()
self.stopped = False
self.dataloder = dataloder
self.batchSize = batchSize
# initialize the queue used to store frames read from
# the video file
self.Q = LifoQueue(maxsize=queueSize)
def __init__(self):
self.scales = "1,2,3"
self.batch_size = 1
self.confidence = 0.5
self.nms_thesh = 0.4
self.reso = 416
self.CUDA = False
self.num_classes = 80
self.classes = load_classes('data/coco.names')
self.colors = load_colors('data/pallete')
self.model = Darknet('cfg/yolov3.cfg', self.reso)
self.model.load_state_dict(torch.load('yolov3.pkl'))
self.inp_dim = self.reso
assert self.inp_dim % 32 == 0
assert self.inp_dim > 32
if self.CUDA:
self.model.cuda()
self.model.eval()
def set_yolo(args):
labelsPath = os.path.sep.join([args["yolo"], "coco.names"])
labels = load_classes(labelsPath)
weightsPath = os.path.sep.join([args["yolo"], "yolov3.weights"])
configPath = os.path.sep.join([args["yolo"], "yolov3.cfg"])
# load our YOLO object detector trained on COCO dataset (80 classes)
# and determine only the *output* layer names that we need from YOLO
print("[INFO] loading YOLO from disk...")
model = Darknet(configPath)
model.load_weights(weightsPath)
model.net_info["height"] = 320
model.cuda()
model.eval()
return labels, model
def __init__(self, webcam=0, batchSize=1, queueSize=256):
# initialize the file video stream along with the boolean
# used to indicate if the thread should be stopped or not
self.det_model = Darknet("yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights('models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda()
self.det_model.eval()
self.stream = cv2.VideoCapture(int(webcam))
assert self.stream.isOpened(), 'Cannot open webcam'
self.stopped = False
self.batchSize = batchSize
# initialize the queue used to store frames read from
# the video file
self.Q = LifoQueue(maxsize=queueSize)
def parse_net(layer: Dict[str, str]) -> Darknet:
steps = [int(s.strip()) for s in layer['steps'].split(',')]
scales = [float(s.strip()) for s in layer['scales'].split(',')]
steps_scales = list(zip(steps, scales))
return Darknet(int(layer['batch']), int(layer['subdivisions']),
int(layer['max_batches']), int(layer['width']),
int(layer['height']), int(layer['channels']),
float(layer['learning_rate']), int(layer['burn_in']),
steps_scales, float(layer['momentum']),
float(layer['decay']), float(layer['angle']),
float(layer['saturation']), float(layer['exposure']),
float(layer['hue']))
def __init__(self, dataloder, batchSize=1, queueSize=1024):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# initialize the file video stream along with the boolean
# used to indicate if the thread should be stopped or not
self.det_model = Darknet(
"/Users/yunyi/Desktop/AlphaPose/yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights(
'/Users/yunyi/Desktop/AlphaPose/models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.to(device)
self.det_model.eval()
self.stopped = False
self.dataloder = dataloder
self.batchSize = batchSize
# initialize the queue used to store frames read from
# the video file
self.Q = LifoQueue(maxsize=queueSize)
def __init__(self, dataloder, batchSize=1, queueSize=1024):
self.det_model = Darknet("yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights('models/yolo/yolov3-spp.weights')
self.det_model.net_info['height'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['height'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda()
self.det_model.eval()
self.stopped = False
self.dataloder = dataloder
self.batchSize = batchSize
# initialize the queue used to store frames read from
# the video file
self.Q = LifoQueue(maxsize=queueSize)
pose_dataset = Mscoco()
if opt.fast_inference:
self.pose_model = InferenNet_fast(4 * 1 + 1, pose_dataset)
else:
self.pose_model = InferenNet(4 * 1 + 1, pose_dataset)
self.pose_model.cuda()
self.pose_model.eval()
def __init__(self):
# set args
self.args = self.arg_parse()
self.confidence = float(self.args.confidence)
self.nms_thesh = float(self.args.nms_thresh)
# load file
self.cfgfile = "yolo/cfg/yolov3.cfg"
self.weightsfile = "yolo/yolov3.weights"
self.classes = load_classes('yolo/data/coco.names')
self.colors = pkl.load(open("yolo/pallete", "rb"))
# set model
self.num_classes = 80
self.bbox_attrs = 5 + self.num_classes
self.model = Darknet(self.cfgfile)
self.model.load_weights(self.weightsfile)
self.model.net_info["height"] = self.args.reso
self.inp_dim = int(self.model.net_info["height"])
self.model.eval()
def __init__(self, dataloder, batchSize=1, queueSize=1024):
self.det_model = Darknet("yolo/cfg/yolov3-spp.cfg")
self.det_model.load_weights("models/yolo/yolov3-spp.weights")
self.det_model.net_info['hight'] = opt.inp_dim
self.det_inp_dim = int(self.det_model.net_info['hight'])
assert self.det_inp_dim % 32 == 0
assert self.det_inp_dim > 32
self.det_model.cuda()
self.det_model.eval()
self.stopped = False
self.dataloder = dataloder
self.batchSize = batchSize
self.datalen = self.dataloder.length()
leftover = 0
if (self.datalen) % batchSize:
leftover = 1
self.num_batches = self.datalen // batchSize + leftover
if opt.sp:
self.Q = Queue(maxsize=queueSize)
else:
self.Q = mp.Queue(maxsize=queueSize)
