from utils.tools import load_class_names, plot_boxes_cv2
from utils.detect_tools import *
device = "cuda" if torch.cuda.is_available() else "cpu"
if __name__ == '__main__':
width = 416
height = 416
namesfile = 'data/coco.names'
model = Yolov4(inference=True).to(device)
static_dict = torch.load("weights/yolov4.pth")
model.load_state_dict(static_dict)
cap = cv2.VideoCapture(0)
ret, image = cap.read()
class_names = load_class_names(namesfile)
for i in range(1000):
ret, img = cap.read()
sized = cv2.resize(img, (width, height))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
boxes = pytorch_detect(model, sized, 0.4, 0.6, device)
image = plot_boxes_cv2(img, boxes[0], class_names=class_names)
cv2.imshow("result", image)
cv2.waitKey(10)
print(np.mean(np.array(infer_times[1:])))
torch.set_default_tensor_type('torch.FloatTensor')
# YEAR = '2007'
devkit_path = VOCroot
dataset_mean = (128.0, 128.0, 128.0)
set_type = 'test'
num_classes = len(VOC_CLASSES) + 1 # +1 background
conf_thres = 0.5
nms_thresh = 0.25
input_size = 300
ovthresh = 0.4
top_k = 100
label_file = "/home/hyer/workspace/algo/Detection/SSD/ssd.pytorch/data/k1.txt"
# weightfile = '/home/hyer/workspace/algo/Detection/SSD/ssd.pytorch/models/mobv1-ssd-k1/T16184/VOC_mobile_300_19105321_10000.pth'
class_names = load_class_names(label_file)
version = config.VOC_mobile_300_19105321
trained_model = '/home/hyer/workspace/algo/Detection/SSD/ssd.pytorch/models/mobv1-ssd-k1/add_0/VOC_mobile_300_19105321_40000.pth'
net = build_ssd('test',
version,
size=300,
num_classes=num_classes,
bkg_label=0,
top_k=top_k,
conf_thresh=conf_thres,
nms_thresh=nms_thresh) # initialize SSD
net.load_weights(trained_model)
net.eval()
print('==> Finished loading model!')
# # load data