File size: 1,402 Bytes
ab26132 f8a743c ab26132 f8a743c ab26132 f8a743c ab26132 f8a743c ab26132 f8a743c ab26132 f8a743c ab26132 f8a743c |
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 |
from PIL import Image
import base64
import io
import numpy as np
# Здесь предполагается, что у вас есть функция segment(image: PIL.Image) -> np.ndarray (маска)
from medsam2_model import MedSAM2
# Загрузка модели
model = MedSAM2("MedSAM2_pretrain_10ep_b1_AMD-SD_sam2_hiera_t.pth")
def image_to_base64(image: Image.Image) -> str:
buffered = io.BytesIO()
image.save(buffered, format="PNG")
return "data:image/png;base64," + base64.b64encode(buffered.getvalue()).decode()
def predict(payload):
# CVAT будет отправлять изображение в виде base64
if isinstance(payload, dict) and "image" in payload:
image_data = payload["image"]
if image_data.startswith("data:image"):
header, base64_data = image_data.split(",", 1)
image = Image.open(io.BytesIO(base64.b64decode(base64_data)))
# Получаем маску
mask_array = model.predict(image_np, box) # Предполагается бинарная маска (0 и 1)
mask_pil = Image.fromarray((mask_array * 255).astype(np.uint8))
return [{
"label": "mock-segmentation",
"mask": image_to_base64(mask_pil),
"score": 0.99
}]
return [{"label": "mock-segmentation", "mask": None, "score": 0.0}]
|