INNER CODE UNIT · Python
get_segmentator
davidefiocco/streamlit-fastapi-model-serving · fastapi/segmentation.py:10
def get_segmentator():
model = torch.hub.load(
"pytorch/vision:v0.6.0", "deeplabv3_resnet101", pretrained=True
)
model.eval()
return model
def get_segments(model, binary_image, max_size=512):
input_image = Image.open(io.BytesIO(binary_image)).convert("RGB")
width, height = input_image.size
resize_factor = min(max_size / width, max_size / height)
resized_image = input_image.resize(
(
int(input_image.width * resize_factor),