def main(img_path, base_name, checkpoint_path):
ori_imgs, framed_imgs, framed_metas = preprocess(img_path, max_size=input_size)
if use_cuda:
x = torch.stack([torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
x = x.to(torch.float32 if not use_float16 else torch.float16).permute(0, 3, 1, 2)
model = EfficientDetBackbone(compound_coef=compound_coef, num_classes=len(obj_list),
ratios=anchor_ratios, scales=anchor_scales)
# model.load_state_dict(torch.load(f'weights/efficientdet-d{compound_coef}.pth'))
model.load_state_dict(torch.load(checkpoint_path))
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
with torch.no_grad():
features, regression, classification, anchors = model(x)
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
out = postprocess(x,
anchors, regression, classification,
regressBoxes, clipBoxes,
threshold, iou_threshold)
out = invert_affine(framed_metas, out)
display(out, ori_imgs, base_name,imshow=False, imwrite=True)
def load_model(compound_coef, obj_list, params, weights_path, use_cuda,
use_float16):
model = EfficientDetBackbone(compound_coef=compound_coef,
num_classes=len(obj_list),
ratios=eval(params['anchors_ratios']),
scales=eval(params['anchors_scales']))
model.load_state_dict(
torch.load(weights_path, map_location=torch.device('cpu')))
model.requires_grad_(False)
model.eval()
if use_cuda:
model.cuda(gpu)
if use_float16:
model.half()
return model
def read_images():
for filename in os.listdir(imgfile_path):
ori_imgs, framed_imgs, framed_metas = preprocess(os.path.join(
imgfile_path, filename),
max_size=input_size)
if use_cuda:
x = torch.stack(
[torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
x = x.to(torch.float32 if not use_float16 else torch.float16).permute(
0, 3, 1, 2)
model = EfficientDetBackbone(compound_coef=7,
num_classes=len(obj_list),
ratios=anchor_ratios,
scales=anchor_scales)
model.load_state_dict(
torch.load(f'weights/efficientdet-d7/efficientdet-d7.pth')
) #place weight path here
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
with torch.no_grad():
features, regression, classification, anchors = model(x)
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
out = postprocess(x, anchors, regression, classification,
regressBoxes, clipBoxes, threshold,
iou_threshold)
out = invert_affine(framed_metas, out)
display(filename, out, ori_imgs, imshow=False, imwrite=True)
print('running speed test...')
with torch.no_grad():
print('test1: model inferring and postprocessing')
print('inferring image for 10 times...')
t1 = time.time()
for _ in range(10):
_, regression, classification, anchors = model(x)
out = postprocess(x, anchors, regression, classification,
regressBoxes, clipBoxes, threshold,
iou_threshold)
out = invert_affine(framed_metas, out)
t2 = time.time()
tact_time = (t2 - t1) / 10
print(f'{tact_time} seconds, {1 / tact_time} FPS, @batch_size 1')
def test(threshold=0.2):
with open("datasets/vcoco/new_prior_mask.pkl", "rb") as file:
prior_mask = pickle.load(file, encoding="bytes")
model = EfficientDetBackbone(num_classes=len(eval(params["obj_list"])),
num_union_classes=25,
num_inst_classes=51,
compound_coef=args.compound_coef,
ratios=eval(params["anchors_ratios"]),
scales=eval(params["anchors_scales"]))
model.load_state_dict(
torch.load(weights_path, map_location=torch.device('cpu')))
model.requires_grad_(False)
model.eval()
if args.cuda:
model = model.cuda()
if args.float16:
model = model.half()
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
img_dir = os.path.join(data_dir, "vcoco/coco/images/%s" % "val2014")
with open(os.path.join(data_dir, 'vcoco/data/splits/vcoco_test.ids'),
'r') as f:
image_ids = f.readlines()
image_ids = [int(id) for id in image_ids]
_t = {'im_detect': Timer(), 'misc': Timer()}
detection = []
for i, image_id in enumerate(image_ids):
_t['im_detect'].tic()
file = "COCO_val2014_" + (str(image_id)).zfill(12) + '.jpg'
img_detection = img_detect(file,
img_dir,
model,
input_size,
regressBoxes,
clipBoxes,
prior_mask,
threshold=threshold)
detection.extend(img_detection)
if need_visual:
visual(img_detection, image_id)
_t['im_detect'].toc()
print('im_detect: {:d}/{:d}, average time: {:.3f}s'.format(
i + 1, len(image_ids), _t['im_detect'].average_time))
with open(detection_path, "wb") as file:
pickle.dump(detection, file)
def test(threshold=0.2):
model = EfficientDetBackbone(num_classes=num_objects,
num_union_classes=num_union_actions,
num_inst_classes=num_inst_actions,
compound_coef=args.compound_coef,
ratios=eval(params["anchors_ratios"]),
scales=eval(params["anchors_scales"]))
model.load_state_dict(
torch.load(weights_path, map_location=torch.device('cpu')))
model.requires_grad_(False)
model.eval()
if args.cuda:
model = model.cuda()
if args.float16:
model = model.half()
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
img_dir = os.path.join(data_dir,
"hico_20160224_det/images/%s" % "test2015")
_t = {'im_detect': Timer(), 'misc': Timer()}
detection = {}
count = 0
for line in glob.iglob(img_dir + '/' + '*.jpg'):
count += 1
_t['im_detect'].tic()
image_id = int(line[-9:-4])
file = "HICO_test2015_" + (str(image_id)).zfill(8) + ".jpg"
# if file != "COCO_val2014_000000001987.jpg":
# continue
dets = img_detect(file,
img_dir,
model,
input_size,
regressBoxes,
clipBoxes,
threshold=threshold)
detection[image_id] = dets
# detection.extend(img_detection)
_t['im_detect'].toc()
print('im_detect: {:d}/{:d}, average time: {:.3f}s'.format(
count, 9658, _t['im_detect'].average_time))
with open(detection_path, "wb") as file:
pickle.dump(detection, file)
def test(opt):
params = Params(f'projects/{opt.project}.yml')
project_name = params.project_name
obj_list = params.obj_list
compound_coef = opt.compound_coef
force_input_size = None # set None to use default size
img_dir = opt.img_dir
model_path = opt.model_path
use_cuda = True
use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True
input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536]
input_size = input_sizes[compound_coef] if force_input_size is None else force_input_size
model = EfficientDetBackbone(compound_coef=compound_coef, num_classes=len(obj_list))
model.load_state_dict(torch.load(model_path))
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
gt = COCO(opt.ann_file)
gt_lst = load_coco_bboxes(gt, is_gt=True)
imgs = glob.glob(os.path.join(img_dir, '*.jpg'))
det_lst = []
progressbar = tqdm(imgs)
for i, img in enumerate(progressbar):
det = single_img_test(img, input_size, model, use_cuda, use_float16)
det_lst.extend(det)
progressbar.update()
progressbar.set_description('Step: {}/{}'.format(i, len(imgs)))
evaluator = Evaluator()
ret, mAP = evaluator.GetMAPbyClass(
gt_lst,
det_lst,
method='EveryPointInterpolation'
)
# Get metric values per each class
for metricsPerClass in ret:
cl = metricsPerClass['class']
ap = metricsPerClass['AP']
ap_str = '{0:.3f}'.format(ap)
print('AP: %s (%s)' % (ap_str, cl))
mAP_str = '{0:.3f}'.format(mAP)
print('mAP: %s\n' % mAP_str)
def model_fn(model_dir):
# based entirely off of
# https://github.com/zylo117/Yet-Another-EfficientDet-Pytorch/blob/master/coco_eval.py
print(f'building and loading efficientdet d{EFFICIENTDET_COMPOUND_COEF}')
model = EfficientDetBackbone(compound_coef=EFFICIENTDET_COMPOUND_COEF,
num_classes=len(PARAMS['obj_list']),
ratios=eval(PARAMS['anchors_ratios']),
scales=eval(PARAMS['anchors_scales']))
state_dict = torch.hub.load_state_dict_from_url(
url=get_weights_url(c=EFFICIENTDET_COMPOUND_COEF),
model_dir=model_dir,
map_location=torch.device('cpu'))
model.load_state_dict(state_dict)
model.requires_grad_(False)
model.eval()
if USE_CUDA:
model.cuda(0)
if USE_FLOAT16:
model.half()
return model
def eval(pretrained_weights: Path, inputs_splitted_into_lists: list,
compound_coef: int, use_cuda: bool) -> list:
threshold = 0.2
iou_threshold = 0.2
# replace this part with your project's anchor config
anchor_ratios = [(1.0, 1.0), (1.4, 0.7), (0.7, 1.4)]
anchor_scales = [2**0, 2**(1.0 / 3.0), 2**(2.0 / 3.0)]
model = EfficientDetBackbone(compound_coef=compound_coef,
num_classes=1,
ratios=anchor_ratios,
scales=anchor_scales)
model.load_state_dict(torch.load(pretrained_weights, map_location='cpu'))
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
predictions = []
for inputs_split in inputs_splitted_into_lists:
with torch.no_grad():
features, regression, classification, anchors = model(inputs_split)
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
out = postprocess(inputs_split, anchors, regression,
classification, regressBoxes, clipBoxes,
threshold, iou_threshold)
predictions += out
return predictions
class EfficientDet(object):
obj_list = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
'train', 'truck', 'boat', 'traffic light',
'fire hydrant', '', 'stop sign', 'parking meter', 'bench',
'bird', 'cat', 'dog', 'horse', 'sheep',
'cow', 'elephant', 'bear', 'zebra', 'giraffe', '', 'backpack',
'umbrella', '', '', 'handbag', 'tie',
'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
'kite', 'baseball bat', 'baseball glove',
'skateboard', 'surfboard', 'tennis racket', 'bottle', '',
'wine glass', 'cup', 'fork', 'knife', 'spoon',
'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli',
'carrot', 'hot dog', 'pizza', 'donut',
'cake', 'chair', 'couch', 'potted plant', 'bed', '',
'dining table', '', '', 'toilet', '', 'tv',
'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink',
'refrigerator', '', 'book', 'clock', 'vase', 'scissors',
'teddy bear', 'hair drier', 'toothbrush']
def __init__(self, weightfile, score_thresh,
nms_thresh, is_xywh=True, use_cuda=True, use_float16=False):
print('Loading weights from %s... Done!' % (weightfile))
# constants
self.score_thresh = score_thresh
self.nms_thresh = nms_thresh
self.use_cuda = use_cuda
self.is_xywh = is_xywh
compound_coef = 0
force_input_size = None # set None to use default size
self.use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True
# tf bilinear interpolation is different from any other's, just make do
input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536]
self.input_size = input_sizes[compound_coef] if \
force_input_size is None else force_input_size
# load model
self.model = EfficientDetBackbone(compound_coef=compound_coef,
num_classes=len(self.obj_list))
# f'weights/efficientdet-d{compound_coef}.pth'
self.model.load_state_dict(torch.load(weightfile))
self.model.requires_grad_(False)
self.model.eval()
if self.use_cuda:
self.model = self.model.cuda()
if self.use_float16:
self.model = self.model.half()
# Box
self.regressBoxes = BBoxTransform()
self.clipBoxes = ClipBoxes()
def __call__(self, imgs):
# frame preprocessing
_, framed_imgs, framed_metas = preprocess(imgs,
max_size=self.input_size)
if self.use_cuda:
x = torch.stack([torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
dtype = torch.float32 if not self.use_float16 else torch.float16
x = x.to(dtype).permute(0, 3, 1, 2)
# model predict
with torch.no_grad():
features, regression, classification, anchors = self.model(x)
out = postprocess(x,
anchors, regression, classification,
self.regressBoxes, self.clipBoxes,
self.score_thresh, self.nms_thresh)
# result
out = invert_affine(framed_metas, out)
if len(out) == 0:
return None, None, None
rois = [o['rois'] for o in out]
scores = [o['scores'] for o in out]
class_ids = [o['class_ids'] for o in out]
if self.is_xywh:
return xyxy_to_xywh(rois), scores, class_ids
else:
return rois, scores, class_ids
def effdet_detection(content, effdet):
video_src = 0 # set int to use webcam, set str to read from a video file
compound_coef = 0
force_input_size = None # set None to use default size
threshold = 0.5
iou_threshold = 0.2
use_cuda = True
use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True
obj_list = [
'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train',
'truck', 'boat', 'traffic light', 'fire hydrant', '', 'stop sign',
'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep',
'cow', 'elephant', 'bear', 'zebra', 'giraffe', '', 'backpack',
'umbrella', '', '', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
'skateboard', 'surfboard', 'tennis racket', 'bottle', '', 'wine glass',
'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
'chair', 'couch', 'potted plant', 'bed', '', 'dining table', '', '',
'toilet', '', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator',
'', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
'toothbrush'
]
# tf bilinear interpolation is different from any other's, just make do
input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536]
input_size = input_sizes[
compound_coef] if force_input_size is None else force_input_size
# load model
model = EfficientDetBackbone(compound_coef=compound_coef,
num_classes=len(obj_list))
model.load_state_dict(
torch.load(f'weights/efficientdet-d{compound_coef}.pth'))
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
# function for display
def display(preds, imgs, content, effdet):
for i in range(len(imgs)):
if len(preds[i]['rois']) == 0:
return imgs[i]
for j in range(len(preds[i]['rois'])):
(x1, y1, x2, y2) = preds[i]['rois'][j].astype(np.int)
#cv2.rectangle(imgs[i], (x1, y1), (x2, y2), (255, 255, 0), 2)
obj = obj_list[preds[i]['class_ids'][j]]
score = float(preds[i]['scores'][j])
if obj == content:
effdet.send_message_to_scratch(
(x1 + x2) * 0.5 * 0.625 - 200) #发送指定类别的识别框位置到scratch
print((x1 + x2) * 0.5 * 0.625 - 200)
cv2.rectangle(imgs[i], (x1, y1), (x2, y2), (255, 255, 0),
2)
cv2.putText(imgs[i], '{}, {:.3f}'.format(obj, score),
(x1, y1 + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
(255, 255, 0), 1)
return imgs[i]
# Box
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
# Video capture
cap = cv2.VideoCapture(video_src)
while True:
ret, frame = cap.read()
if not ret:
break
# frame preprocessing
ori_imgs, framed_imgs, framed_metas = preprocess_video(
frame, max_size=input_size)
if use_cuda:
x = torch.stack(
[torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
x = x.to(torch.float32 if not use_float16 else torch.float16).permute(
0, 3, 1, 2)
# model predict
with torch.no_grad():
features, regression, classification, anchors = model(x)
out = postprocess(x, anchors, regression, classification,
regressBoxes, clipBoxes, threshold,
iou_threshold)
# result
out = invert_affine(framed_metas, out)
img_show = display(out, ori_imgs, content, effdet)
# show frame by frame
cv2.imshow('frame', img_show)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
x = torch.stack([torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
x = x.to(torch.float32 if not use_float16 else torch.float16).permute(0, 3, 1, 2)
model = EfficientDetBackbone(compound_coef=compound_coef, num_classes=len(obj_list),
ratios=anchor_ratios, scales=anchor_scales)
model.load_state_dict(torch.load(f'weights/efficientdet-d{compound_coef}.pth'))
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
with torch.no_grad():
features, regression, classification, anchors = model(x)
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
out = postprocess(x,
anchors, regression, classification,
regressBoxes, clipBoxes,
threshold, iou_threshold)
def display(preds, imgs, imshow=True, imwrite=False):
for i in range(len(imgs)):
def EfficientDetNode():
rospy.init_node('efficient_det_node', anonymous=True)
rospy.Subscriber('input', String, image_callback, queue_size=1)
pub = rospy.Publisher('/image_detections', Detection2DArray, queue_size=10)
rate = rospy.Rate(1) # 10hz
path_list = os.listdir(path)
path_list.sort(key=lambda x: int(x.split('.')[0]))
stamp_file = open(stamp_path)
stamp_lines = stamp_file.readlines()
stamp_i = 0
for filename in path_list:
img_path = filename
cur_frame = img_path[:-4]
img_path = path + "/" + img_path
cur_stamp = ((float)(stamp_lines[stamp_i][-13:].strip('\n')))
# cur_stamp = rospy.Time.from_sec(
# ((float)(stamp_lines[stamp_i][-13:].strip('\n'))))
stamp_i += 1
detection_results = Detection2DArray()
# tf bilinear interpolation is different from any other's, just make do
input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536, 1536]
input_size = input_sizes[
compound_coef] if force_input_size is None else force_input_size
ori_imgs, framed_imgs, framed_metas = preprocess(img_path,
max_size=input_size)
if use_cuda:
x = torch.stack(
[torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
x = x.to(torch.float32 if not use_float16 else torch.float16).permute(
0, 3, 1, 2)
model = EfficientDetBackbone(compound_coef=compound_coef,
num_classes=len(obj_list),
ratios=anchor_ratios,
scales=anchor_scales)
model.load_state_dict(
torch.load(f'weights/efficientdet-d{compound_coef}.pth',
map_location='cpu'))
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
with torch.no_grad():
features, regression, classification, anchors = model(x)
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
out = postprocess(x, anchors, regression, classification,
regressBoxes, clipBoxes, threshold,
iou_threshold)
out = invert_affine(framed_metas, out)
display(cur_frame, out, ori_imgs, imshow=False, imwrite=True)
for i in range(len(out)):
for j in range(len(out[i]['rois'])):
x1, y1, x2, y2 = out[i]['rois'][j].astype(np.int)
obj = obj_list[out[i]['class_ids'][j]]
score = float(out[i]['scores'][j])
result = ObjectHypothesisWithPose()
result.score = score
if (obj == 'car'):
result.id = 0
if (obj == 'person'):
result.id = 1
if (obj == 'cyclist'):
result.id = 2
detection_msg = Detection2D()
detection_msg.bbox.center.x = (x1 + x2) / 2
detection_msg.bbox.center.y = (y1 + y2) / 2
detection_msg.bbox.size_x = x2 - x1
detection_msg.bbox.size_y = y2 - y1
detection_msg.results.append(result)
detection_results.detections.append(detection_msg)
rospy.loginfo("%d: %lf", detection_msg.results[0].id,
detection_msg.results[0].score)
detection_results.header.seq = cur_frame
#detection_results.header.stamp = cur_stamp
rospy.loginfo(detection_results.header.stamp)
pub.publish(detection_results)
if not os.path.exists(txt_path):
os.makedirs(txt_path)
#with open(f'txt/{cur_frame}.txt', 'w') as f:
with open(f'{txt_path}/{cur_frame}.txt', 'w') as f:
#f.write(str((float)(stamp_lines[stamp_i][-13:].strip('\n'))) + "\n")
f.write(str(cur_stamp) + "\n")
for detection in detection_results.detections:
f.write(str(detection.bbox.center.x) + " ")
f.write(str(detection.bbox.center.y) + " ")
f.write(str(detection.bbox.size_x) + " ")
f.write(str(detection.bbox.size_y) + " ")
f.write(str(detection.results[0].id) + " ")
f.write(str(detection.results[0].score) + "\n")
f.close()
rate.sleep()
print('running speed test...')
with torch.no_grad():
print('test1: model inferring and postprocessing')
print('inferring image for 10 times...')
t1 = time.time()
for _ in range(10):
_, regression, classification, anchors = model(x)
out = postprocess(x, anchors, regression, classification,
regressBoxes, clipBoxes, threshold,
iou_threshold)
out = invert_affine(framed_metas, out)
t2 = time.time()
tact_time = (t2 - t1) / 10
print(f'{tact_time} seconds, {1 / tact_time} FPS, @batch_size 1')
def test(opt):
compound_coef = 2
force_input_size = None # set None to use default size
img_id = opt.img_id
img_path = opt.img_path
img_path = img_path + str(img_id) + '.jpg'
# replace this part with your project's anchor config
anchor_ratios = [(1.0, 1.0), (1.4, 0.7), (0.7, 1.4)]
anchor_scales = [2**0, 2**(1.0 / 3.0), 2**(2.0 / 3.0)]
threshold = 0.2
iou_threshold = 0.2
use_cuda = True
use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True
obj_list = ['02010001', '02010002']
color_list = standard_to_bgr(STANDARD_COLORS)
input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536, 1536]
input_size = input_sizes[
compound_coef] if force_input_size is None else force_input_size
ori_imgs, framed_imgs, framed_metas = preprocess(img_path,
max_size=input_size)
if use_cuda:
x = torch.stack([torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
x = x.to(torch.float32 if not use_float16 else torch.float16).permute(
0, 3, 1, 2)
model = EfficientDetBackbone(compound_coef=compound_coef,
num_classes=len(obj_list),
ratios=anchor_ratios,
scales=anchor_scales)
model.load_state_dict(torch.load(opt.weights, map_location='cpu'))
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
with torch.no_grad():
features, regression, classification, anchors = model(x)
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
out = postprocess(x, anchors, regression, classification, regressBoxes,
clipBoxes, threshold, iou_threshold)
def display(preds, imgs, imshow=True, imwrite=False, img_id=1):
for i in range(len(imgs)):
if len(preds[i]['rois']) == 0:
continue
imgs[i] = imgs[i].copy()
imgs[i] = cv2.cvtColor(imgs[i], cv2.COLOR_BGR2RGB)
for j in range(len(preds[i]['rois'])):
x1, y1, x2, y2 = preds[i]['rois'][j].astype(np.int)
obj = obj_list[preds[i]['class_ids'][j]]
score = float(preds[i]['scores'][j])
plot_one_box(imgs[i], [x1, y1, x2, y2],
label=obj,
score=score,
color=color_list[get_index_label(obj, obj_list)])
if imshow:
cv2.imshow('img', imgs[i])
cv2.waitKey(0)
if imwrite:
str1 = 'test/' + str(img_id) + '.jpg'
cv2.imwrite(str1, imgs[i])
out = invert_affine(framed_metas, out)
display(out, ori_imgs, imshow=False, imwrite=True, img_id=img_id)
print('running speed test...')
with torch.no_grad():
print('test1: model inferring and postprocessing')
print('inferring image for 10 times...')
t1 = time.time()
for _ in range(10):
_, regression, classification, anchors = model(x)
out = postprocess(x, anchors, regression, classification,
regressBoxes, clipBoxes, threshold,
iou_threshold)
out = invert_affine(framed_metas, out)
tempList = []
for j in range(len(out[0]['class_ids'])):
tempout = {}
tempout['image_id'] = img_id
if out[0]['class_ids'][j] == 1:
tempout['category_id'] = 2
else:
tempout['category_id'] = 1
tempout['score'] = out[0]['scores'][j].astype(np.float64)
tempout['bbox'] = [
(out[0]['rois'][j][0]).astype(np.float64),
(out[0]['rois'][j][1]).astype(np.float64),
(out[0]['rois'][j][2]).astype(np.float64) -
(out[0]['rois'][j][0]).astype(np.float64),
(out[0]['rois'][j][3]).astype(np.float64) -
(out[0]['rois'][j][1]).astype(np.float64),
]
tempList.append(tempout)
t2 = time.time()
tact_time = (t2 - t1) / 10
print(f'{tact_time} seconds, {1 / tact_time} FPS, @batch_size 1')
with open("test/" + str(img_id) + ".json", "w") as f:
json.dump(tempList, f)
print("生成标注后的图片(" + str(img_id) + ".jpg)和json(" + str(img_id) +
".json)到test文件夹中...")
model_2.load_state_dict(
torch.load(
f'/data/efdet/logs/{project}/crop/weights/{save_time2}/efficientdet-d{compound_coef}_{number}.pth',
map_location='cpu'))
model_1.requires_grad_(False)
model_1.eval()
model_2.requires_grad_(False)
model_2.eval()
if use_cuda:
model_1 = model_1.cuda()
model_2 = model_2.cuda()
if use_float16:
model_1 = model_1.half()
model_2 = model_2.half()
def display(out_1, out_2, imgs, imshow=True, showtime=0, imwrite=False):
# if len(preds[i]['rois']) == 0: # if model dosen't detect object, not show image
# continue
for img, out_1 in zip(imgs, out_1):
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
for i in range(len(out_1['rois'])):
ox1, oy1, ox2, oy2 = out_1['rois'][i].astype(np.int)
obj_1 = obj_list_1[out_1['class_ids'][i]]
score = float(out_1['scores'][i])
color = color_list[get_index_label(obj_1, obj_list_1)]
plot_one_box(img, [ox1, oy1, ox2, oy2],
def infer(self, image):
img = np.array(image)
img = img[:, :, ::-1] #rgb 2 bgr
anchor_ratios = [(1.0, 1.0), (1.4, 0.7), (0.7, 1.4)]
anchor_scales = [2**0, 2**(1.0 / 3.0), 2**(2.0 / 3.0)]
threshold = 0.25
iou_threshold = 0.25
force_input_size = None
use_cuda = False
use_float16 = False
cudnn.fastest = False
cudnn.benchmark = False
input_size = 512
ori_imgs, framed_imgs, framed_metas = preprocess(img,
max_size=input_size)
if use_cuda:
x = torch.stack(
[torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
x = x.to(torch.float32 if not use_float16 else torch.float16).permute(
0, 3, 1, 2)
model = EfficientDetBackbone(compound_coef=0,
num_classes=len(self.labels),
ratios=anchor_ratios,
scales=anchor_scales)
model.load_state_dict(torch.load(self.path, map_location='cpu'))
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
with torch.no_grad():
features, regression, classification, anchors = model(x)
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
out = postprocess(x, anchors, regression, classification,
regressBoxes, clipBoxes, threshold,
iou_threshold)
pred = invert_affine(framed_metas, out)
results = []
for i in range(len(ori_imgs)):
if len(pred[i]['rois']) == 0:
continue
ori_imgs[i] = ori_imgs[i].copy()
for j in range(len(pred[i]['rois'])):
xt1, yt1, xbr, ybr = pred[i]['rois'][j].astype(np.float64)
xt1 = float(xt1)
yt1 = float(yt1)
xbr = float(xbr)
yb4 = float(ybr)
obj = str(pred[i]['class_ids'][j])
obj_label = self.labels.get(obj)
obj_score = str(pred[i]['scores'][j])
results.append({
"confidence": str(obj_score),
"label": obj_label,
"points": [xt1, yt1, xbr, ybr],
"type": "rectangle",
})
return results
class Model():
def __init__(self, compound_coef=0, force_input_size=512, threshold=0.2, iou_threshold=0.2):
self.compound_coef = compound_coef
self.force_input_size = force_input_size # set None to use default size
self.threshold = threshold
self.iou_threshold = iou_threshold
self.use_cuda = True
self.use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True
self.obj_list = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', '', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep',
'cow', 'elephant', 'bear', 'zebra', 'giraffe', '', 'backpack', 'umbrella', '', '', 'handbag', 'tie',
'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
'skateboard', 'surfboard', 'tennis racket', 'bottle', '', 'wine glass', 'cup', 'fork', 'knife', 'spoon',
'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut',
'cake', 'chair', 'couch', 'potted plant', 'bed', '', 'dining table', '', '', 'toilet', '', 'tv',
'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
'refrigerator', '', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
'toothbrush']
# tf bilinear interpolation is different from any other's, just make do
self.input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536]
self.input_size = self.input_sizes[self.compound_coef] if self.force_input_size is None else self.force_input_size
self.model = EfficientDetBackbone(
compound_coef=self.compound_coef, num_classes=len(self.obj_list))
self.model.load_state_dict(torch.load(
f'weights/efficientdet-d{self.compound_coef}.pth'))
self.model.requires_grad_(False)
self.model.eval()
if self.use_cuda:
self.model = self.model.cuda()
if self.use_float16:
self.model = self.model.half()
def predict(self, raw_img):
self.ori_imgs, self.framed_imgs, self.framed_metas = preprocess_raw(raw_img, max_size=self.input_size)
if self.use_cuda:
x = torch.stack([torch.from_numpy(fi).cuda() for fi in self.framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in self.framed_imgs], 0)
x = x.to(torch.float32 if not self.use_float16 else torch.float16).permute(0, 3, 1, 2)
with torch.no_grad():
self.features, self.regression, self.classification, self.anchors = self.model(x)
self.regressBoxes = BBoxTransform()
self.clipBoxes = ClipBoxes()
out = postprocess(x,
self.anchors, self.regression, self.classification,
self.regressBoxes, self.clipBoxes,
self.threshold, self.iou_threshold)
pred = invert_affine(self.framed_metas, out)
return pred
def label_img(self, preds, imgs):
for i in range(len(imgs)):
if len(preds[i]['rois']) == 0:
continue
for j in range(len(preds[i]['rois'])):
(x1, y1, x2, y2) = preds[i]['rois'][j].astype(np.int)
cv2.rectangle(imgs[i], (x1, y1), (x2, y2), (255, 255, 0), 2)
obj = self.obj_list[preds[i]['class_ids'][j]]
score = float(preds[i]['scores'][j])
cv2.putText(imgs[i], '{}, {:.3f}'.format(obj, score),
(x1, y1 + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
(255, 255, 0), 1)
return imgs
def run(self, raw_img):
pred_label = self.predict(raw_img)
pred_img = self.label_img(pred_label, self.ori_imgs)
return pred_img[0]
def efficientDet_video_inference(video_src,compound_coef = 0,force_input_size=None,
frame_skipping = 3,
threshold=0.2,out_path=None,imshow=False,
display_fps=False):
#deep-sort variables
# Definition of the parameters
max_cosine_distance = 0.3
nn_budget = None
nms_max_overlap = 1.0
model_filename = '/home/shaheryar/Desktop/Projects/Football-Monitoring/deep_sort/model_weights/mars-small128.pb'
encoder = gdet.create_box_encoder(model_filename, batch_size=1)
metric = nn_matching.NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
tracker = Tracker(metric,n_init=5)
# efficientDet-pytorch variables
iou_threshold = 0.4
use_cuda = True
use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True
input_size = input_sizes[compound_coef] if force_input_size is None else force_input_size
# load model
model = EfficientDetBackbone(compound_coef=compound_coef, num_classes=len(obj_list))
model.load_state_dict(torch.load(f'weights/efficientdet-d{compound_coef}.pth'))
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
# Video capture
cap = cv2.VideoCapture(video_src)
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))
fourcc = cv2.VideoWriter_fourcc(*'MPEG')
fps = cap.get(cv2.CAP_PROP_FPS)
print("Video fps",fps)
if(out_path is not None):
outp = cv2.VideoWriter(out_path, fourcc, fps, (frame_width, frame_height))
i=0
start= time.time()
current_frame_fps=0
while True:
ret, frame = cap.read()
if not ret:
break
t1=time.time()
if (frame_skipping==0 or i%frame_skipping==0):
# if(True):
# frame preprocessing (running detections)
ori_imgs, framed_imgs, framed_metas, t1 = preprocess_video(frame, width=input_size, height=input_size)
if use_cuda:
x = torch.stack([fi.cuda() for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
# model predict
t1=time.time()
with torch.no_grad():
features, regression, classification, anchors = model(x)
out = postprocess(x,
anchors, regression, classification,
regressBoxes, clipBoxes,
threshold, iou_threshold)
# Post processing
out = invert_affine(framed_metas, out)
# decoding bbox ,object name and scores
boxes,classes,scores =decode_predictions(out[0])
org_boxes = boxes.copy()
t2 = time.time() - t1
# feature extraction for deep sort
boxes = [convert_bbox_to_deep_sort_format(frame.shape, b) for b in boxes]
features = encoder(frame,boxes)
detections = [Detection(bbox, 1.0, feature) for bbox, feature in zip(boxes, features)]
boxes = np.array([d.tlwh for d in detections])
# print(boxes)
scores = np.array([d.confidence for d in detections])
indices = preprocessing.non_max_suppression(boxes, nms_max_overlap, scores)
detections = [detections[i] for i in indices]
tracker.predict()
tracker.update(detections)
i = i + 1
img_show=frame.copy()
for j in range(len(org_boxes)):
img_show =drawBoxes(img_show,org_boxes[j],(255,255,0),str(tracker.tracks[j].track_id))
for track in tracker.tracks:
if not track.is_confirmed() or track.time_since_update > 1:
continue
bbox = track.to_tlbr()
x1=int(bbox[0])
y1 = int(bbox[1])
x2 = int(bbox[2])
y2=int(bbox[3])
roi= frame[y1:y2,x1:x2]
cv2.rectangle(img_show, (x1, y1), (x2, y2), update_color_association(roi, track.track_id), 2)
cv2.putText(img_show, str(track.track_id), (x1, y1), 0, 5e-3 * 100, (255, 255, 0), 1)
if display_fps:
current_frame_fps=1/t2
else:
current_frame_fps=0
cv2.putText(img_show, 'FPS: {0:.2f}'.format(current_frame_fps), (30, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
(255, 255, 0),
2, cv2.LINE_AA)
if (i % int(fps) == 0):
print("Processed ", str(int(i / fps)), "seconds")
print("Time taken",time.time()-start)
# print(color_dict)
if imshow:
img_show=cv2.resize(img_show,(0,0),fx=0.75,fy=0.75)
cv2.imshow('Frame',img_show)
# Press Q on keyboard to exit
if cv2.waitKey(1) & 0xFF == ord('q'):
break
if out_path is not None:
outp.write(img_show)
cap.release()
outp.release()
def excuteModel(videoname):
# Video's path
# set int to use webcam, set str to read from a video file
if videoname is not None:
video_src = os.path.join(r'D:\GitHub\Detection\server\uploads', f"{videoname}.mp4")
else:
video_src = 'D:\\GitHub\\Detection\\server\AImodel\\videotest\\default.mp4'
compound_coef = 2
trained_weights = 'D:\\GitHub\\Detection\\server\\AImodel\\weights\\efficientdet-video.pth'
force_input_size = None # set None to use default size
threshold = 0.2
iou_threshold = 0.2
use_cuda = True
use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True
obj_list = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', '', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep',
'cow', 'elephant', 'bear', 'zebra', 'giraffe', '', 'backpack', 'umbrella', '', '', 'handbag', 'tie',
'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
'skateboard', 'surfboard', 'tennis racket', 'bottle', '', 'wine glass', 'cup', 'fork', 'knife', 'spoon',
'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut',
'cake', 'chair', 'couch', 'potted plant', 'bed', '', 'dining table', '', '', 'toilet', '', 'tv',
'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
'refrigerator', '', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
'toothbrush']
# tf bilinear interpolation is different from any other's, just make do
input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536]
input_size = input_sizes[compound_coef] if force_input_size is None else force_input_size
# load model
model = EfficientDetBackbone(
compound_coef=compound_coef, num_classes=len(obj_list))
model.load_state_dict(torch.load(trained_weights))
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
# function for display
# Box
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
# Video capture
cap = cv2.VideoCapture(video_src)
length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
writer = None
# try to determine the total number of frames in the video file
try:
prop = cv2.cv.CV_CAP_PROP_FRAME_COUNT if imutils.is_cv2() \
else cv2.CAP_PROP_FRAME_COUNT
total = int(vs.get(prop))
print("[INFO] {} total frames in video".format(total))
# an error occurred while trying to determine the total
# number of frames in the video file
except:
print("[INFO] could not determine # of frames in video")
total = -1
path_out = os.path.join(os.path.dirname(
os.path.abspath(__file__)), 'outvideo')
path_result = r"D:\GitHub\Detection\server\AImodel\videotest\default.mp4"
path_asset = r"D:\GitHub\Detection\client\src\assets"
for i in range(0, length):
ret, frame = cap.read()
if not ret:
break
# frame preprocessing
ori_imgs, framed_imgs, framed_metas = preprocess_video(
frame, max_size=input_size)
if use_cuda:
x = torch.stack([torch.from_numpy(fi).cuda()
for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
x = x.to(torch.float32 if not use_float16 else torch.float16).permute(
0, 3, 1, 2)
# model predict
with torch.no_grad():
features, regression, classification, anchors = model(x)
out = postprocess(x,
anchors, regression, classification,
regressBoxes, clipBoxes,
threshold, iou_threshold)
# result
out = invert_affine(framed_metas, out)
img_show = display(out, ori_imgs, obj_list)
if writer is None:
# initialize our video writer
fourcc = 0x00000021
#fourcc = cv2.VideoWriter_fourcc(*'mp4v')
if videoname is not None:
path_result = os.path.join(path_out, f"{videoname}.mp4")
else:
path_result = os.path.join(path_out, "default.mp4")
writer = cv2.VideoWriter(path_result, fourcc, 30, (img_show.shape[1], img_show.shape[0]), True)
# write the output frame to disk
writer.write(img_show)
print("Processing data... " + str(round((i+1)/length, 3)*100) + " %")
# show frame by frame
#cv2.imshow('frame', img_show)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
print("[INFO] cleaning up...")
writer.release()
cap.release()
cv2.destroyAllWindows()
if videoname is not None:
path_asset = os.path.join(path_asset, f"{videoname}.mp4")
else:
path_asset = os.path.join(path_asset, "default.mp4")
copyfile(path_result, path_asset)
return path_asset
def getImageDetections(imagePath, weights, nms_threshold, confidenceParam, coefficient):
"""
Runs the detections and returns all detection into a single structure.
Parameters
----------
imagePath : str
Path to all images.
weights : str
path to the weights.
nms_threshold : float
non-maximum supression threshold.
confidenceParam : float
confidence score for the detections (everything above this threshold is considered a valid detection).
coefficient : int
coefficient of the current efficientdet model (from d1 to d7).
Returns
-------
detectionsList : List
return a list with all predicted bounding-boxes.
"""
compound_coef = coefficient
force_input_size = None # set None to use default size
img_path = imagePath
threshold = confidenceParam
iou_threshold = nms_threshold
use_cuda = True
use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True
obj_list = ['class_name']
# tf bilinear interpolation is different from any other's, just make do
input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536]
input_size = input_sizes[compound_coef] if force_input_size is None else force_input_size
ori_imgs, framed_imgs, framed_metas = preprocess(img_path, max_size=input_size)
if use_cuda:
x = torch.stack([torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)
else:
x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)
x = x.to(torch.float32 if not use_float16 else torch.float16).permute(0, 3, 1, 2)
model = EfficientDetBackbone(compound_coef=compound_coef, num_classes=len(obj_list),
# replace this part with your project's anchor config
ratios=[(1.0, 1.0), (1.4, 0.7), (0.7, 1.4)],
scales=[2 ** 0, 2 ** (1.0 / 3.0), 2 ** (2.0 / 3.0)])
model.load_state_dict(torch.load(rootDir+'logs/' + project + '/' + weights))
model.requires_grad_(False)
model.eval()
if use_cuda:
model = model.cuda()
if use_float16:
model = model.half()
with torch.no_grad():
features, regression, classification, anchors = model(x)
regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()
out = postprocess(x,
anchors, regression, classification,
regressBoxes, clipBoxes,
threshold, iou_threshold)
out = invert_affine(framed_metas, out)
for i in range(len(ori_imgs)):
if len(out[i]['rois']) == 0:
continue
detectionsList = []
for j in range(len(out[i]['rois'])):
(x1, y1, x2, y2) = out[i]['rois'][j].astype(np.int)
detectionsList.append((float(out[i]['scores'][j]), x1, y1, x2, y2))
return detectionsList
coco_eval.summarize()
if __name__ == '__main__':
SET_NAME = params['val_set']
VAL_GT = f'datasets/{params["project_name"]}/{SET_NAME}.json'
VAL_IMGS = f'datasets/{params["project_name"]}/{SET_NAME}/{SET_NAME}'
MAX_IMAGES = 10000
coco_gt = COCO(VAL_GT)
image_ids = coco_gt.getImgIds()[:MAX_IMAGES]
if override_prev_results or not os.path.exists(
f'{SET_NAME}_bbox_results.json'):
model = EfficientDetBackbone(compound_coef=compound_coef,
num_classes=len(obj_list),
ratios=eval(params['anchors_ratios']),
scales=eval(params['anchors_scales']))
model.load_state_dict(torch.load(weights_path))
model.requires_grad_(False)
model.eval()
if use_cuda:
model.cuda(gpu)
if use_float16:
model.half()
evaluate_coco(VAL_IMGS, SET_NAME, image_ids, coco_gt, model)
# _eval(coco_gt, image_ids, f'{SET_NAME}_bbox_results.json')
def __init__(self,
video_src: str,
video_output: str,
text_output: str,
obj_list: list,
input_sizes: list,
reid_cpkt: str,
compound_coef: int,
force_input_size=None,
threshold=0.2,
iou_threshold=0.2,
use_cuda=True,
use_float16=False,
cudnn_fastest=True,
cudnn_benchmark=True,
max_dist=0.2,
min_confidence=0.3,
nms_max_overlap=0.5,
max_iou_distance=0.7,
max_age=70,
n_init=3,
nn_budget=100,
selected_target=None):
# I/O
# Video's path
self.video_src = video_src # set int to use webcam, set str to read from a video file
self.video_output = video_output # output to the specific position
# text path
self.text_output = text_output # output to the file with the csv format
# DETECTOR
self.compound_coef = compound_coef
self.force_input_size = force_input_size # set None to use default size
self.threshold = threshold
self.iou_threshold = iou_threshold
self.use_cuda = use_cuda
self.use_float16 = use_float16
cudnn.fastest = cudnn_fastest
cudnn.benchmark = cudnn_benchmark
# coco_name
self.obj_list = obj_list
# input size
self.input_sizes = input_sizes
self.input_size = input_sizes[self.compound_coef] if force_input_size is None else force_input_size
# load detector model
model = EfficientDetBackbone(compound_coef=self.compound_coef, num_classes=len(obj_list))
model.load_state_dict(torch.load(f'weights/efficientdet-d{self.compound_coef}.pth'))
model.requires_grad_(False)
model.eval()
if self.use_cuda and torch.cuda.is_available():
self.detector = model.cuda()
if self.use_float16:
self.detector = model.half()
# TRACKER
self.reid_cpkt = reid_cpkt
self.max_dist = max_dist
self.min_confidence = min_confidence
self.nms_max_overlap = nms_max_overlap
self.max_iou_distance = max_iou_distance
self.max_age = max_age
self.n_init = n_init
self.nn_budget = nn_budget
# load tracker model,
self.trackers = []
self.selected_target = selected_target
for num in range(0, len(self.selected_target)):
self.trackers.append(build_tracker(reid_cpkt,
max_dist,
min_confidence,
nms_max_overlap,
max_iou_distance,
max_age,
n_init,
nn_budget,
use_cuda))
# video frames
self.frame_id = 0
