15 lines
426 B
Python
15 lines
426 B
Python
from pathlib import Path
|
|
|
|
from app.modules.weights.service import classify_weight
|
|
|
|
|
|
def test_classify_weight():
|
|
item = classify_weight(Path("Seg_All_In_One_YoloModel/yolov8n-seg.pt"))
|
|
assert item["family"] == "yolo"
|
|
assert item["role"] == "weight"
|
|
|
|
best = classify_weight(Path("Seg_Predict_YoloModel/YOLOv9e-seg/weights/best.pt"))
|
|
assert best["family"] == "yolo"
|
|
assert best["role"] == "trained_best"
|
|
|