File size: 2,845 Bytes
06a2085
ec500f1
 
 
 
 
 
06a2085
095f7cc
ec500f1
 
 
 
c973a85
06a2085
ec500f1
 
 
 
1d08b5b
 
ec500f1
 
 
06a2085
ec500f1
1d08b5b
 
ec500f1
1d08b5b
 
 
 
c973a85
 
 
06a2085
 
 
 
 
 
 
 
 
 
 
c973a85
06a2085
 
 
 
 
 
ec500f1
 
 
 
 
1d08b5b
 
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 coco_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...")
    coco_anns = []

    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.extend(yolo_boxes_to_coco_annotations(image_id, result.boxes,
                                                            confidence_threshold=confidence_threshold))

    evaluator.update(coco_anns)
    evaluator.synchronize_between_processes()
    evaluator.accumulate()
    evaluator.summarize()


def yolo_evaluator(model: YoloModel, dataset: YoloDataset):
    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.material.__enter__()
            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": int(image_id),
            "category_id": int(box.cls.tolist()[0]),
            "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()