File size: 1,041 Bytes
8d9a780 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | 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):
# Load all your gingivitis models here
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):
# Convert image to OpenCV format
img = np.array(image)
# Run detection (choose one for now)
results = self.model_redness(img)
# Annotate the result image
annotated = results[0].plot() # returns a NumPy array (with boxes)
_, buffer = cv2.imencode(".png", annotated)
img_bytes = buffer.tobytes()
# Return base64 image string
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)
|