predict.py

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)