YOLOv8.py 1.06 KB
Newer Older
1
2
3
4
from ultralytics import YOLO


class YOLOv8MFDModel(object):
5
    def __init__(self, weight, device="cpu"):
6
        self.mfd_model = YOLO(weight)
7
8
        if not device.startswith("cpu"):
            self.mfd_model.half()
9
10
11
        self.device = device

    def predict(self, image):
12
        mfd_res = self.mfd_model.predict(
13
            image, imgsz=1888, conf=0.25, iou=0.45, verbose=False, device=self.device
14
        )[0]
15
16
        return mfd_res

17
18
19
    def batch_predict(self, images: list, batch_size: int) -> list:
        images_mfd_res = []
        for index in range(0, len(images), batch_size):
20
21
22
23
24
25
26
            mfd_res = [
                image_res.cpu()
                for image_res in self.mfd_model.predict(
                    images[index : index + batch_size],
                    imgsz=1888,
                    conf=0.25,
                    iou=0.45,
27
                    verbose=False,
28
29
30
                    device=self.device,
                )
            ]
31
32
33
            for image_res in mfd_res:
                images_mfd_res.append(image_res)
        return images_mfd_res