Spaces:
Runtime error
Runtime error
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 | |