def main():
img_size = 512 # 必须是32的整数倍 [416, 512, 608]
cfg = "cfg/yolov3-spp.cfg"
weights = "weights/yolov3-spp-ultralytics-{}.pt".format(img_size)
img_path = "test.jpg"
input_size = (img_size, img_size)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Darknet(cfg, img_size)
model.load_state_dict(torch.load(weights, map_location=device)["model"])
model.to(device)
model.eval()
# init
img = torch.zeros((1, 3, img_size, img_size), device=device)
model(img)
img_o = cv2.imread(img_path) # BGR
assert img_o is not None, "Image Not Found " + img_path
img = img_utils.letterbox(img_o, new_shape=input_size, auto=True, color=(0, 0, 0))[0]
# Convert
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(device).float()
img /= 255.0 # scale (0, 255) to (0, 1)
img = img.unsqueeze(0) # add batch dimension
t1 = torch_utils.time_synchronized()
pred = model(img)[0] # only get inference result
t2 = torch_utils.time_synchronized()
print(t2 - t1)
pred = utils.non_max_suppression(pred, conf_thres=0.3, iou_thres=0.6, multi_label=True)[0]
t3 = time.time()
print(t3 - t2)
# process detections
pred[:, :4] = utils.scale_coords(img.shape[2:], pred[:, :4], img_o.shape).round()
print(pred.shape)
bboxes = pred[:, :4].detach().cpu().numpy()
scores = pred[:, 4].detach().cpu().numpy()
classes = pred[:, 5].detach().cpu().numpy().astype(np.int) + 1
category_index = dict([(i + 1, str(i + 1)) for i in range(90)])
img_o = draw_box(img_o[:, :, ::-1], bboxes, classes, scores, category_index)
plt.imshow(img_o)
plt.show()
img_o.save("test_result.jpg")
def main():
img_size = 512 # 必须是32的整数倍 [416, 512, 608]
cfg = "cfg/my_yolov3.cfg" # 改成生成的.cfg文件
weights = "weights/yolov3spp-voc-512.pth".format(img_size) # 改成自己训练好的权重文件
json_path = "./data/pascal_voc_classes.json" # json标签文件
img_path = "test.jpg"
assert os.path.exists(cfg), "cfg file {} dose not exist.".format(cfg)
assert os.path.exists(weights), "weights file {} dose not exist.".format(
weights)
assert os.path.exists(json_path), "json file {} dose not exist.".format(
json_path)
assert os.path.exists(img_path), "image file {} dose not exist.".format(
img_path)
json_file = open(json_path, 'r')
class_dict = json.load(json_file)
category_index = {v: k for k, v in class_dict.items()}
input_size = (img_size, img_size)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = YOLOV3_SPP(cfg, img_size)
model.load_state_dict(torch.load(weights, map_location=device)["model"])
model.to(device)
model.eval()
with torch.no_grad():
# init
img = torch.zeros((1, 3, img_size, img_size), device=device)
model(img)
img_o = cv2.imread(img_path) # BGR
assert img_o is not None, "Image Not Found " + img_path
img = img_utils.letterbox(img_o,
new_shape=input_size,
auto=True,
color=(0, 0, 0))[0]
# Convert
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(device).float()
img /= 255.0 # scale (0, 255) to (0, 1)
img = img.unsqueeze(0) # add batch dimension
t1 = torch_utils.time_synchronized()
pred = model(img)[0] # only get inference result
t2 = torch_utils.time_synchronized()
print(t2 - t1)
pred = utils.non_max_suppression(pred,
conf_thres=0.1,
iou_thres=0.6,
multi_label=True)[0]
t3 = time.time()
print(t3 - t2)
# process detections
pred[:, :4] = utils.scale_coordinates(pred[:, :4], img.shape[2:],
img_o.shape).round()
print(pred.shape)
bboxes = pred[:, :4].detach().cpu().numpy()
scores = pred[:, 4].detach().cpu().numpy()
classes = pred[:, 5].detach().cpu().numpy().astype(np.int) + 1
img_o = draw_box(img_o[:, :, ::-1], bboxes, classes, scores,
category_index)
plt.imshow(img_o)
plt.show()
img_o.save("test_result.jpg")
def main():
img_size = 512 # 必须是32的整数倍 [416, 512, 608]
cfg = "cfg/yolov3-spp.cfg"
weights = "weights/yolov3-spp-ultralytics-{}.pt".format(img_size)
assert os.path.exists(cfg), "cfg file does not exist..."
assert os.path.exists(weights), "weights file does not exist..."
input_size = (img_size, img_size) # [h, w]
# create model
model = models.Darknet(cfg, input_size)
# load model weights
model.load_state_dict(torch.load(weights, map_location=device)["model"])
model.to(device)
model.eval()
# input to the model
# [batch, channel, height, width]
# x = torch.rand(1, 3, *input_size, requires_grad=True)
img_path = "test.jpg"
img_o = cv2.imread(img_path) # BGR
assert img_o is not None, "Image Not Found " + img_path
# preprocessing img
img = img_utils.letterbox(img_o,
new_shape=input_size,
auto=False,
color=(0, 0, 0))[0]
# Convert
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img).astype(np.float32)
img /= 255.0 # scale (0, 255) to (0, 1)
img = np.expand_dims(img, axis=0) # add batch dimension
x = torch.tensor(img)
torch_out = model(x)
save_path = "yolov3spp.onnx"
# export the model
torch.onnx.export(
model, # model being run
x, # model input (or a tuple for multiple inputs)
save_path, # where to save the model (can be a file or file-like object)
export_params=
True, # store the trained parameter weights inside the model file
opset_version=12, # the ONNX version to export the model to
do_constant_folding=
True, # whether to execute constant folding for optimization
input_names=["images"], # the model's input names
# output_names=["classes", "boxes"], # the model's output names
output_names=["prediction"],
dynamic_axes={
"images": {
0: "batch_size"
}, # variable length axes
"prediction": {
0: "batch_size"
}
})
# "classes": {0: "batch_size"},
# "confidence": {0: "batch_size"},
# "boxes": {0: "batch_size"}})
# check onnx model
onnx_model = onnx.load(save_path)
onnx.checker.check_model(onnx_model)
# print(onnx.helper.printable_graph(onnx_model.graph))
ort_session = onnxruntime.InferenceSession(save_path)
# compute ONNX Runtime output prediction
ort_inputs = {"images": to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)
# compare ONNX Runtime and Pytorch results
# assert_allclose: Raises an AssertionError if two objects are not equal up to desired tolerance.
np.testing.assert_allclose(to_numpy(torch_out),
ort_outs[0],
rtol=1e-03,
atol=1e-05)
# np.testing.assert_allclose(to_numpy(torch_out[1]), ort_outs[1], rtol=1e-03, atol=1e-05)
# np.testing.assert_allclose(to_numpy(torch_out[2]), ort_outs[2], rtol=1e-03, atol=1e-05)
print(
"Exported model has been tested with ONNXRuntime, and the result looks good!"
)