giladvdn
/

test-sam-handler

Mask Generation

Inference Endpoints

Model card Files Files and versions Community

test-sam-handler / handler.py

Gilad Avidan

change output format

5f41ad3 11 months ago

raw history blame contribute delete

No virus

2.29 kB

	import base64
	import gzip
	import numpy as np
	from io import BytesIO
	from typing import Dict, List, Any
	from PIL import Image
	import torch
	from transformers import SamModel, SamProcessor

	def pack_bits(boolean_tensor):
	# Flatten the tensor and add padding if necessary
	flat = boolean_tensor.flatten()
	if flat.size()[0] % 8 != 0:
	padding = np.zeros((8 - flat.size % 8,), dtype=bool)
	flat = np.concatenate([flat, padding])

	# Reshape into bytes and pack into binary string
	packed = np.packbits(flat.reshape((-1, 8)))
	packed = packed.tobytes()
	return gzip.compress(packed)
	# json_str = json.dumps({"shape": boolean_tensor.shape, "data": binary_str})


	class EndpointHandler():
	def __init__(self, path=""):
	# Preload all the elements you are going to need at inference.
	self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	self.model = SamModel.from_pretrained("facebook/sam-vit-base").to(self.device)
	self.processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

	def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
	"""
	data args:
	inputs (:obj: `str` \| `PIL.Image` \| `np.array`)
	kwargs
	Return:
	A :obj:`list` \| `dict`: will be serialized and returned
	"""

	inputs = data.pop("inputs", data)
	parameters = data.pop("parameters", {"mode": "image"})

	# Decode base64 image to PIL
	image = Image.open(BytesIO(base64.b64decode(inputs['image']))).convert("RGB")
	input_points = [inputs['points']] # 2D localization of a window

	model_inputs = self.processor(image, input_points=input_points, return_tensors="pt").to(self.device)
	outputs = self.model(**model_inputs)
	masks = self.processor.image_processor.post_process_masks(
	outputs.pred_masks.cpu(),
	model_inputs["original_sizes"].cpu(),
	model_inputs["reshaped_input_sizes"].cpu())
	scores = outputs.iou_scores

	packed = [base64.b64encode(pack_bits(masks[0][0][i])).decode() for i in range(masks[0].shape[1])]
	shape = list(masks[0].shape)[2:]

	return {"masks": packed, "scores": scores[0][0].tolist(), "shape": shape}