def demo_cam(cfg, save_path, save):
net = yolo(cfg)
net.load_state_dict(torch.load(cfg.trained_model))
if cfg.cuda: net = net.cuda()
net.eval()
cam = cv2.VideoCapture(0)
if not cam.isOpened():
raise IOError("check your camera or the opencv library...")
if save:
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(save_path + '/out_camera.avi', fourcc, 20.0,
(640, 480))
while cam.isOpened():
ret, image = cam.read()
if ret:
reimg = cv2.resize(image, cfg.image_size)
reimg = torch.from_numpy(reimg.transpose(
2, 0, 1)).float().div(255.0).unsqueeze(0)
reimg = Variable(reimg, volatile=True)
if cfg.cuda: reimg = reimg.cuda()
boxes, scores, classes = net(reimg,
(image.shape[1], image.shape[0]))
boxes, scores, classes = boxes.data.cpu(), scores.data.cpu(
), classes.data.cpu()
for i, c in list(enumerate(classes)):
pred_class, box, score = cfg.classes[c], boxes[i], scores[i]
label = '{} {:.2f}'.format(pred_class, score)
draw_box_cv(cfg, image, label, box, c)
cv2.imshow('camera', image)
if cv2.waitKey(1) & 0xFF == ord('q'): break
if save: out.write(image)
cam.release()
if save: out.release()
cv2.destroyAllWindows()
def demo_video(cfg, file, save_path, save, start_frame=0):
net = yolo(cfg)
net.load_state_dict(torch.load(cfg.trained_model))
if cfg.cuda: net = net.cuda()
net.eval()
video = cv2.VideoCapture(file)
if not video.isOpened():
raise IOError('Could not open video, check your opencv and video')
if save:
fps = video.get(cv2.CAP_PROP_FPS)
size = (int(video.get(cv2.CAP_PROP_FRAME_WIDTH)),
int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)))
out = cv2.VideoWriter(os.path.join(save_path + file.split('/')[-1]),
cv2.VideoWriter_fourcc(*'XVID'), fps, size)
if start_frame > 0:
video.set(cv2.CAP_PROP_POS_MSEC, start_frame)
accum_time = 0
curr_fps = 0
fps = "FPS: "
prev_time = timer()
while True:
retval, orig_image = video.read()
if not retval:
print('Done !!!')
return
reimg = cv2.resize(orig_image, cfg.image_size)
reimg = torch.from_numpy(reimg.transpose(
2, 0, 1)).float().div(255.0).unsqueeze(0)
x = Variable(reimg, volatile=True)
if cfg.cuda: x = x.cuda()
boxes, scores, classes = net(
x, (orig_image.shape[1], orig_image.shape[0]))
boxes, scores, classes = boxes.data.cpu(), scores.data.cpu(
), classes.data.cpu()
for i, c in list(enumerate(classes)):
pred_class, box, score = cfg.classes[c], boxes[i], scores[i]
label = '{} {:.2f}'.format(pred_class, score)
draw_box_cv(cfg, orig_image, label, box, c)
curr_time = timer()
exec_time = curr_time - prev_time
prev_time = curr_time
accum_time += exec_time
curr_fps = curr_fps + 1
if accum_time > 1:
accum_time -= 1
fps = "FPS:" + str(curr_fps)
curr_fps = 0
cv2.rectangle(orig_image, (0, 0), (50, 17), (255, 255, 255), -1)
cv2.putText(orig_image, fps, (3, 10), 0, 0.35, (0, 0, 0), 1)
if save:
out.write(orig_image)
cv2.imshow("yolo result", orig_image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
video.release()
if save:
out.release()
def demo(cfg, img_list, save_path):
transform = transforms.Compose(
[transforms.Resize(cfg.image_size),
transforms.ToTensor()])
net = yolo(cfg)
net.load_state_dict(torch.load(cfg.trained_model))
if cfg.cuda: net = net.cuda()
net.eval()
for img in img_list:
image = Image.open(img)
reimg = Variable(transform(image).unsqueeze(0))
if cfg.cuda: reimg = reimg.cuda()
boxes, scores, classes = net(reimg, image.size)
boxes, scores, classes = boxes.data.cpu(), scores.data.cpu(
), classes.data.cpu()
print('Find {} boxes for {}.'.format(len(boxes), img.split('/')[-1]))
for i, c in list(enumerate(classes)):
pred_class, box, score = cfg.classes[c], boxes[i], scores[i]
label = '{} {:.2f}'.format(pred_class, score)
draw_box(cfg, image, label, box, c)
image.save(os.path.join(save_path, img.split('/')[-1]), quality=90)
def weight2pth(config_path, weights_path, output_path):
assert config_path.endswith('.cfg'), '{} is not a .cfg file'.format(config_path)
assert weights_path.endswith('.weights'), '{} is not a .weights file'.format(weights_path)
# weights header
weights_file = open(weights_path, 'rb')
weights_header = np.ndarray(shape=(4,), dtype='int32', buffer=weights_file.read(16))
print('Weights Header: ', weights_header)
# convert config information
unique_config_file = unique_config_sections(config_path)
cfg_parser = configparser.ConfigParser()
cfg_parser.read_file(unique_config_file)
# network information
cfgname = config_path.split('/')[-1].split('.')[0]
cfg = importlib.import_module('config.' + cfgname)
net = yolo(cfg)
net_dict = net.state_dict()
keys = list(net_dict.keys())
key_num, count, prev_filter = 0, 0, 3
print('loading the weights ...')
for section in cfg_parser.sections():
if section.startswith('convolutional'):
filters = int(cfg_parser[section]['filters'])
size = int(cfg_parser[section]['size'])
bn = 'batch_normalize' in cfg_parser[section]
activation = cfg_parser[section]['activation']
# three special case
if section == 'convolutional_20':
prev_filter = 512
elif section == 'convolutional_21':
prev_filter = 1280
elif section == 'convolutional_0':
prev_filter = 3
else:
prev_filter = weights_shape[0]
weights_shape = (filters, prev_filter, size, size)
weights_size = np.product(weights_shape)
print('conv2d', 'bn' if bn else ' ', activation, weights_shape)
conv_bias = np.ndarray(
shape=(filters,),
dtype='float32',
buffer=weights_file.read(filters * 4))
count += filters
if bn:
bn_weights = np.ndarray(
shape=(3, filters),
dtype='float32',
buffer=weights_file.read(filters * 12))
count += 3 * filters
net_dict[keys[key_num + 1]].copy_(torch.from_numpy(bn_weights[0]))
net_dict[keys[key_num + 2]].copy_(torch.from_numpy(conv_bias))
net_dict[keys[key_num + 3]].copy_(torch.from_numpy(bn_weights[1]))
net_dict[keys[key_num + 4]].copy_(torch.from_numpy(bn_weights[2]))
else:
net_dict[keys[key_num + 1]].copy_(torch.from_numpy(conv_bias))
# 导入卷积层参数
conv_weights = np.ndarray(
shape=weights_shape,
dtype='float32',
buffer=weights_file.read(weights_size * 4))
count += weights_size
net_dict[keys[key_num]].copy_(torch.from_numpy(conv_weights))
key_num = key_num + 5 if bn else key_num + 1
else:
continue
# check the convert
remaining_weights = len(weights_file.read()) // 4
weights_file.close()
print('Read {} of {} from Darknet weights.'.format(count, count + remaining_weights))
if remaining_weights > 0:
print('Warning: {} unused weights'.format(remaining_weights))
# save the net.state_dict
torch.save(net_dict, os.path.join(output_path, cfgname + '.pth'))