def test_load_from_yolov5_torchscript(arch, size_divisible, version, upstream_version, hash_prefix):
import cv2
from yolort.utils import read_image_to_tensor
from yolort.v5 import letterbox
# Loading and pre-processing the image
img_path = "test/assets/zidane.jpg"
img_raw = cv2.imread(img_path)
img = letterbox(img_raw, new_shape=(640, 640), stride=size_divisible)[0]
img = read_image_to_tensor(img)
base_url = "https://github.com/ultralytics/yolov5/releases/download/"
model_url = f"{base_url}/{upstream_version}/{arch}.pt"
checkpoint_path = attempt_download(model_url, hash_prefix=hash_prefix)
score_thresh = 0.25
model = YOLO.load_from_yolov5(checkpoint_path, score_thresh=score_thresh, version=version)
model.eval()
scripted_model = torch.jit.script(model)
scripted_model.eval()
out = model(img[None])
out_script = scripted_model(img[None])
torch.testing.assert_close(out[0]["scores"], out_script[1][0]["scores"], rtol=0, atol=0)
torch.testing.assert_close(out[0]["labels"], out_script[1][0]["labels"], rtol=0, atol=0)
torch.testing.assert_close(out[0]["boxes"], out_script[1][0]["boxes"], rtol=0, atol=0)
def test_load_from_yolov5(arch, size_divisible, version, upstream_version, hash_prefix):
img_path = "test/assets/bus.jpg"
base_url = "https://github.com/ultralytics/yolov5/releases/download/"
model_url = f"{base_url}/{upstream_version}/{arch}.pt"
checkpoint_path = attempt_download(model_url, hash_prefix=hash_prefix)
score_thresh = 0.25
model_yolov5 = YOLOv5.load_from_yolov5(
checkpoint_path,
score_thresh=score_thresh,
version=version,
size_divisible=size_divisible,
)
model_yolov5.eval()
out_from_yolov5 = model_yolov5.predict(img_path)
assert isinstance(out_from_yolov5[0], dict)
assert isinstance(out_from_yolov5[0]["boxes"], Tensor)
assert isinstance(out_from_yolov5[0]["labels"], Tensor)
assert isinstance(out_from_yolov5[0]["scores"], Tensor)
model = models.__dict__[arch](
upstream_version=version,
pretrained=True,
score_thresh=score_thresh,
)
model.eval()
out = model.predict(img_path)
torch.testing.assert_close(out_from_yolov5[0]["scores"], out[0]["scores"], rtol=0, atol=0)
torch.testing.assert_close(out_from_yolov5[0]["labels"], out[0]["labels"], rtol=0, atol=0)
torch.testing.assert_close(out_from_yolov5[0]["boxes"], out[0]["boxes"], rtol=0, atol=0)
def test_attempt_download():
model_url = "https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5s.pt"
checkpoint_path = attempt_download(model_url, hash_prefix="9ca9a642")
with open(checkpoint_path, "rb") as f:
bytes = f.read() # read entire file as bytes
readable_hash = hashlib.sha256(bytes).hexdigest()
assert readable_hash[:8] == "9ca9a642"
def test_yolo_graphsurgeon_wo_nms(arch, version, upstream_version, hash_prefix):
base_url = "https://github.com/ultralytics/yolov5/releases/download/"
model_url = f"{base_url}/{upstream_version}/{arch}.pt"
checkpoint_path = attempt_download(model_url, hash_prefix=hash_prefix)
yolo_gs = YOLOTRTGraphSurgeon(checkpoint_path, version=version, enable_dynamic=False)
onnx_file_path = f"yolo_graphsurgeon_wo_nms_{arch}_{hash_prefix}.onnx"
assert not Path(onnx_file_path).exists()
yolo_gs.save(onnx_file_path)
assert Path(onnx_file_path).exists()
def test_yolo_trt_inference_to_onnx(arch, version, upstream_version, hash_prefix):
base_url = "https://github.com/ultralytics/yolov5/releases/download/"
model_url = f"{base_url}/{upstream_version}/{arch}.pt"
checkpoint_path = attempt_download(model_url, hash_prefix=hash_prefix)
model = YOLOTRTInference(checkpoint_path, version=version)
model.eval()
onnx_file_path = f"yolo_trt_inference_to_onnx_{arch}_{hash_prefix}.onnx"
assert not Path(onnx_file_path).exists()
model.to_onnx(onnx_file_path)
assert Path(onnx_file_path).exists()
def test_load_yolov5_model():
img_path = "test/assets/zidane.jpg"
model_url = "https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5s.pt"
checkpoint_path = attempt_download(model_url, hash_prefix="9ca9a642")
model = load_yolov5_model(checkpoint_path, autoshape=True, verbose=False)
results = model(img_path)
assert isinstance(results.pred, list)
assert len(results.pred) == 1
assert isinstance(results.pred[0], Tensor)
assert results.pred[0].shape == (3, 6)
def test_yolo_trt_inference(arch, version, upstream_version, hash_prefix):
base_url = "https://github.com/ultralytics/yolov5/releases/download/"
model_url = f"{base_url}/{upstream_version}/{arch}.pt"
checkpoint_path = attempt_download(model_url, hash_prefix=hash_prefix)
model = YOLOTRTInference(checkpoint_path, version=version)
model.eval()
samples = torch.rand(1, 3, 320, 320)
outs = model(samples)
assert isinstance(outs, tuple)
assert len(outs) == 2
assert isinstance(outs[0], Tensor)
assert isinstance(outs[1], Tensor)
def test_load_from_ultralytics_voc(
arch: str,
version: str,
upstream_version: str,
hash_prefix: str,
):
img_path = "test/assets/bus.jpg"
base_url = "https://github.com/ultralytics/yolov5/releases/download/"
model_url = f"{base_url}/{upstream_version}/{arch}.pt"
checkpoint_path = attempt_download(model_url, hash_prefix=hash_prefix)
# Preprocess
img_raw = cv2.imread(img_path)
img = letterbox(img_raw, new_shape=(320, 320))[0]
img = read_image_to_tensor(img)
conf = 0.25
iou = 0.45
# Define YOLOv5 model
model_yolov5 = load_yolov5_model(checkpoint_path)
model_yolov5.conf = conf # confidence threshold (0-1)
model_yolov5.iou = iou # NMS IoU threshold (0-1)
model_yolov5.eval()
with torch.no_grad():
outs = model_yolov5(img[None])[0]
outs = non_max_suppression(outs, conf, iou, agnostic=True)
out_from_yolov5 = outs[0]
# Define yolort model
model_yolort = YOLO.load_from_yolov5(
checkpoint_path,
score_thresh=conf,
version=version,
)
model_yolort.eval()
with torch.no_grad():
out_from_yolort = model_yolort(img[None])
torch.testing.assert_allclose(out_from_yolort[0]["boxes"],
out_from_yolov5[:, :4])
torch.testing.assert_allclose(out_from_yolort[0]["scores"],
out_from_yolov5[:, 4])
torch.testing.assert_allclose(out_from_yolort[0]["labels"],
out_from_yolov5[:, 5].to(dtype=torch.int64))
def test_load_from_ultralytics(
arch: str,
version: str,
upstream_version: str,
hash_prefix: str,
use_p6: bool,
):
base_url = "https://github.com/ultralytics/yolov5/releases/download/"
model_url = f"{base_url}/{upstream_version}/{arch}.pt"
checkpoint_path = attempt_download(model_url, hash_prefix=hash_prefix)
model_info = load_from_ultralytics(checkpoint_path, version=version)
assert isinstance(model_info, dict)
assert model_info["num_classes"] == 80
assert model_info["size"] == arch.replace("yolov5", "")
assert model_info["use_p6"] == use_p6
assert len(model_info["strides"]) == 4 if use_p6 else 3