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 | |