| from ultralytics import YOLO |
| from PIL import Image |
| import torch |
| import io |
| import base64 |
| import numpy as np |
| import cv2 |
| from huggingface_hub import InferenceMixin |
|
|
| class Model: |
| def __init__(self): |
| |
| self.model_redness = YOLO("best_redness2.pt") |
| self.model_swelling = YOLO("best_swelling.pt") |
| self.model_bleeding = YOLO("bleeding2.pt") |
|
|
| def predict(self, image: Image.Image): |
| |
| img = np.array(image) |
|
|
| |
| results = self.model_redness(img) |
|
|
| |
| annotated = results[0].plot() |
| _, buffer = cv2.imencode(".png", annotated) |
| img_bytes = buffer.tobytes() |
|
|
| |
| img_base64 = base64.b64encode(img_bytes).decode("utf-8") |
| return {"image_base64": img_base64} |
|
|
| model = Model() |
|
|
| def predict(image: Image.Image): |
| return model.predict(image) |
|
|
