You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
76 lines
1.6 KiB
76 lines
1.6 KiB
import torch |
|
|
|
from ultralytics import YOLO |
|
|
|
|
|
def test_model_init(): |
|
model = YOLO.new("yolov8n.yaml") |
|
model.info() |
|
try: |
|
YOLO() |
|
except Exception: |
|
print("Successfully caught constructor assert!") |
|
raise Exception("constructor error didn't occur") |
|
|
|
|
|
def test_model_forward(): |
|
model = YOLO.new("yolov8n.yaml") |
|
img = torch.rand(512 * 512 * 3).view(1, 3, 512, 512) |
|
model.forward(img) |
|
model(img) |
|
|
|
|
|
def test_model_info(): |
|
model = YOLO.new("yolov8n.yaml") |
|
model.info() |
|
model = model.load("best.pt") |
|
model.info(verbose=True) |
|
|
|
|
|
def test_model_fuse(): |
|
model = YOLO.new("yolov8n.yaml") |
|
model.fuse() |
|
model.load("best.pt") |
|
model.fuse() |
|
|
|
|
|
def test_visualize_preds(): |
|
model = YOLO.load("best.pt") |
|
model.predict(source="ultralytics/assets") |
|
|
|
|
|
def test_val(): |
|
model = YOLO.load("best.pt") |
|
model.val(data="coco128.yaml", imgsz=32) |
|
|
|
|
|
def test_model_resume(): |
|
model = YOLO.new("yolov8n.yaml") |
|
model.train(epochs=1, imgsz=32, data="coco128.yaml") |
|
try: |
|
model.resume(task="detect") |
|
except AssertionError: |
|
print("Successfully caught resume assert!") |
|
|
|
|
|
def test_model_train_pretrained(): |
|
model = YOLO.load("best.pt") |
|
model.train(data="coco128.yaml", epochs=1, imgsz=32) |
|
model = model.new("yolov8n.yaml") |
|
model.train(data="coco128.yaml", epochs=1, imgsz=32) |
|
img = torch.rand(512 * 512 * 3).view(1, 3, 512, 512) |
|
model(img) |
|
|
|
|
|
def test(): |
|
test_model_forward() |
|
test_model_info() |
|
test_model_fuse() |
|
test_visualize_preds() |
|
test_val() |
|
test_model_resume() |
|
test_model_train_pretrained() |
|
|
|
|
|
if __name__ == "__main__": |
|
test()
|
|
|