MaskCut

Build error

MaskCut / predict.py

Use docker

5e6d15e over 3 years ago

793 Bytes

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)

	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)