| 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 = package.get('model') | |
| masks = model(image, threshold, num_objects) | |
| for mask in masks: | |
| fg = mask > 0.5 | |
| rgb[fg] = 255 | |
| return Image.fromarray(rgb) |