farkmu45's picture
Add chart
afa9738
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