File size: 793 Bytes
7857874 5e6d15e 7857874 | 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 | import base64
from io import BytesIO
from PIL import Image
import numpy as np
from model import Model
def predict(package, image_base64: str, threshold: float, num_objects: int):
# Decode the image from base64 to PIL.Image
# We use BytesIO to convert the base64 to bytes
base64_split = image_base64.split(',')[1]
buf = BytesIO(base64.b64decode(base64_split))
image = Image.open(buf)
# Get the image path from tmp_image
canvas = Image.new('RGB', image.size, (0, 0, 0))
# We copy the image that and fill it with black, to get the dimensions
rgb = np.array(canvas)
model : Model = package.get('model')
masks = model(image, threshold, num_objects)
for mask in masks:
fg = mask > 0.5
rgb[fg] = 255
return Image.fromarray(rgb) |