def main():
try:
FLAG = process_config()
except:
print("missing or invalid arguments")
exit(0)
if FLAG.GPU_options:
session_config = tf.ConfigProto()
session_config.gpu_options.per_process_gpu_memory_fraction = 0.9
session_config.gpu_options.allow_growth = True
sess = tf.Session(config=session_config)
else:
sess = tf.Session()
model = yolov3(FLAG)
model.build()
model.init_saver()
model.load(sess)
image_test = Image.open('images/timg.jpg')
resized_image = image_test.resize(size=(416, 416))
image_data = np.array(resized_image, dtype='float32') / 255.0
img_hw = tf.placeholder(dtype=tf.float32, shape=[2])
boxes, scores, classes = model.pedict(img_hw,
iou_threshold=0.5,
score_threshold=0.5)
begin_time = time.time()
boxes_, scores_, classes_, conv0 = sess.run(
[boxes, scores, classes, model.feature_extractor.conv0],
feed_dict={
img_hw: [image_test.size[1], image_test.size[0]],
model.x: [image_data]
})
end_time = time.time()
print(end_time - begin_time)
# print conv0
image_draw = draw_boxes(np.array(image_test, dtype=np.float32) / 255,
boxes_,
classes_,
FLAG.names,
scores=scores_)
fig = plt.figure(frameon=False)
ax = plt.Axes(fig, [0, 0, 1, 1])
ax.set_axis_off()
fig.add_axes(ax)
plt.imshow(image_draw)
fig.savefig('prediction.jpg')
plt.show()
sess.close()
def __init__(self):
self.device = torch.device(cfg.device)
self.max_epoch = cfg.max_epoch
self.train_dataloader = dataloader()
self.len_train_dataset = self.train_dataloader.num_annotations
self.model = yolov3().to(self.device)
self.optimizer = torch.optim.SGD(self.model.parameters(),
lr=cfg.lr_start,
momentum=cfg.momentum,
weight_decay=cfg.weight_decay)
self.scheduler = adjust_lr(self.optimizer,
self.max_epoch * self.len_train_dataset,
cfg.lr_start, cfg.lr_end, cfg.warmup)
self.writer = SummaryWriter(cfg.tensorboard_path)
self.iter = 0
return assign_ops
if __name__ == '__main__':
try:
FLAG = process_config()
except:
print("missing or invalid arguments")
exit(0)
if FLAG.npz:
if os.path.exists(FLAG.npz_path):
print "darknet53.conv.74.npz already exists"
else:
print FLAG.config_path
load_weights_for_finetune(FLAG)
elif FLAG.ckpt:
detections = yolov3(FLAG)
# with tf.variable_scope('detector'):
detections.build()
load_ops = load_weights(tf.global_variables(), FLAG.weights_path)
detections.init_saver()
with tf.Session() as sess:
sess.run(load_ops)
writer = tf.summary.FileWriter("experiments/summary/",
graph=sess.graph)
writer.close()
detections.saver.save(sess, "experiments/ckpt/yolov3.ckpt")
else:
raise ValueError('Missing important parameters')
raise FileExistsError("%s is exists" % video_out_path)
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
cap = cv.VideoCapture(video_path)
out = cv.VideoWriter(video_out_path, cv.VideoWriter_fourcc(*"mp4v"),
int(cap.get(cv.CAP_PROP_FPS)),
(int(cap.get(cv.CAP_PROP_FRAME_WIDTH)),
int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))))
assert cap.isOpened()
fps = cap.get(cv.CAP_PROP_FPS)
frame_num = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
print("视频帧率: %s" % fps)
model = yolov3(True).to(device)
model.eval()
for i in range(frame_num):
# 一帧一帧的捕获 跳着读
ret, image_o = cap.read() # BGR
if ret is False:
break # 未读到
image = cv_to_pil(image_o) # -> PIL.Image
image = trans.ToTensor()(image).to(
device) # -> tensor. shape(3, H, W), 0-1
start = time.time()
with torch.no_grad():
target = model(
[image], # image_size=max(image.shape),
score_thresh=score_thresh,
nms_thresh=nms_thresh)[0]
def __init__(self):
self.model = yolov3().to(torch.device(cfg.device))
self.load_weights()
def build_yolo_v3():
model = yolov3().to(device)
model.load_darknet_weights("./checkpoint/darknet53_448.weights")
return model