Processor.py

import PIL.Image as Image
import torch
from fastai.vision.all import Learner
from numpy.typing import NDArray
from typing import List

class Processor():
    def __init__(self, learn: Learner):
        self.__inference = learn
        self.__model = torch.hub.load(
            'ultralytics/yolov5', 'yolov5x6', trust_repo=True)

    def classify_images(self, images: List[NDArray]) -> List[str]:
        result = []
        class_names = self.__inference.dls.vocab

        test_dl = self.__inference.dls.test_dl(images)
        tensors = self.__inference.get_preds(dl=test_dl, with_decoded=True)[2]
        preds = [int(t.item()) for t in tensors]

        for i in preds:
            result.append(class_names[i])

        return result

    def filter_image(self, image: Image) -> bool:
        results = self.__model(image)
        results = results.pandas().xyxy[0]
        person_detected = 0

        for name in results['name']:
            if name == 'person':
                person_detected += 1

        if person_detected == 0 or person_detected > 1:
            return False

        return True