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_image(self, images: NDArray) -> str:
        return self.inference.predict(images)

    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