Spaces:
Runtime error
Runtime error
File size: 2,932 Bytes
06a2085 ec500f1 06a2085 095f7cc ec500f1 06a2085 ec500f1 06a2085 ec500f1 095f7cc ec500f1 06a2085 ec500f1 095f7cc ec500f1 095f7cc 06a2085 095f7cc 06a2085 095f7cc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import shutil
from typing import Callable
from PIL.Image import Image
from coco_eval import CocoEvaluator
from tqdm import tqdm
from yolo_dataset import YoloDataset, MaterialYoloDataset
from yolo_model import YoloModel
image_loader = Callable[[str], Image]
def evaluate(model: YoloModel, dataset: YoloDataset, confidence_threshold=0.6):
coco_gt = dataset.to_coco()
# initialize evaluator with ground truth (gt)
evaluator = CocoEvaluator(coco_gt=coco_gt, iou_types=["bbox"])
print("Running evaluation...")
for image_id, annotations in tqdm(coco_gt.imgToAnns.items()):
# get the inputs
image = coco_gt.imgs[image_id]
results = model.model(source=dataset.load_image(image["file_name"]))
for result in results:
coco_anns = yolo_boxes_to_coco_annotations(image_id, result.boxes,
confidence_threshold=confidence_threshold)
if len(coco_anns) == 0:
continue
evaluator.update(coco_anns)
if len(evaluator.eval_imgs["bbox"]) == 0:
print("No detections!")
return
evaluator.synchronize_between_processes()
evaluator.accumulate()
evaluator.summarize()
return Metrics(model=model, material=dataset.to_material())
class Metrics:
def __init__(self, model: YoloModel, material: MaterialYoloDataset):
self.model = model
self.material = material
self.val = None
def __enter__(self):
if self.val is None:
self.val = self.model.model.val(data=self.material.yaml)
return self.val
def __exit__(self, exc_type, exc_val, exc_tb):
self.material.__exit__(exc_type, exc_val, exc_tb)
shutil.rmtree(self.val.save_dir)
def yolo_boxes_to_coco_annotations(image_id: int, yolo_boxes, confidence_threshold=0.6):
return [
{
"image_id": image_id,
"category_id": box.cls.tolist()[0],
"area": box.xywh.tolist()[0][2] * box.xywh.tolist()[0][3],
"bbox": box.xywh.tolist()[0],
"score": box.conf.tolist()[0],
}
for box in yolo_boxes if box.conf.tolist()[0] > confidence_threshold
]
def main():
yolo_dataset = YoloDataset.from_path('tests/coco8.zip')
coco_gt = yolo_dataset.to_coco()
model = YoloModel("ultralyticsplus/yolov8s", "yolov8s.pt")
# model = YoloModel("SHOU-ISD/fire-and-smoke", "yolov8n.pt")
# evaluate(model=model, coco_gt=coco_gt, loader=yolo_dataset.load_image, confidence_threshold=0.1)
# Validate the model
with yolo_dataset.to_material() as material:
metrics = model.model.val(data=material.yaml)
print(metrics.box.map) # map50-95
print(metrics.box.map50) # map50
print(metrics.box.map75) # map75
print(metrics.box.maps) # a list contains map50-95 of each category
if __name__ == '__main__':
main()
|