Visual segmenter model upload example

models/model_upload/image-segmenter/mask2former-ade/1/model.py

@@ -0,0 +1,139 @@
+import re
+import os
+import sys
+
+sys.path.append(os.path.dirname(__file__))
+from typing import Iterator, List, Tuple
+
+from google.protobuf import json_format
+from clarifai.runners.models.model_runner import ModelRunner
+from clarifai.utils.logging import logger
+from clarifai_grpc.grpc.api import resources_pb2, service_pb2
+from clarifai_grpc.grpc.api.status import status_code_pb2, status_pb2
+from clarifai_grpc.grpc.api import resources_pb2
+
+import yaml
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation 
+
+from utils import *
+
+ROOT = os.path.dirname(__file__)
+
+
+
+class MyRunner(ModelRunner):
+  """A custom runner that adds "Hello World" to the end of the text and replaces the domain of the
+  image URL as an example.
+  """
+
+  def _load_concepts(self, config_path, name, model_path):
+
+    with open(config_path, "r") as f:
+      data = yaml.safe_load(f)
+    if not data.get("concepts"):
+      data = create_concepts_in_yaml(config_path, name, model_path)
+
+    # Map Clarifai concept name to id and reverse
+    self.conceptid2name = {each["id"] : each["name"] for each in data.get("concepts", [])}
+    self.conceptname2id = {each["name"] : each["id"] for each in data.get("concepts", [])}
+
+
+  def load_model(self):
+    """Load the model here."""
+    checkpoint_path = os.path.join(os.path.dirname(__file__), "checkpoints")
+    self.device = 'cuda' #if torch.cuda.is_available() else 'cpu'
+    #self.torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+    logger.info(f"Running on device: {self.device}")
+    self.model = Mask2FormerForUniversalSegmentation.from_pretrained(
+        checkpoint_path, trust_remote_code=True).to(self.device)
+    self.processor = AutoImageProcessor.from_pretrained(checkpoint_path, trust_remote_code=True)
+    self.model.eval()
+    # Load clarifai concept
+    config_path = os.path.join(ROOT, "../config.yaml")
+    self._load_concepts(config_path, "mask2former-ade", checkpoint_path)
+
+    logger.info("Done loading!")
+
+
+  def get_default_infer_kwargs(self, request):
+    infer_kwargs = get_inference_params(request)
+
+    return infer_kwargs
+
+
+  def _model_predict(self, images: List[Image.Image]) -> dict:
+    inputs = self.processor(images=images, return_tensors="pt").to(self.device)
+    with torch.no_grad():
+      outputs = self.model(**inputs)
+    target_sizes = [image.size[::-1] for image in images]
+    results = self.processor.post_process_semantic_segmentation(outputs, target_sizes=target_sizes)
+    outputs = []
+    for i, all_masks_tensor in enumerate(results):
+      masks = {}
+      h, w = target_sizes[i]
+
+      for clss_id in all_masks_tensor.unique().tolist():
+        label = self.model.config.id2label[clss_id]
+        mask = torch.zeros_like(all_masks_tensor)
+        mask[all_masks_tensor == clss_id] = 255
+        mask = mask.cpu().numpy() if self.device == "cuda" else mask.numpy()
+        mask = mask.astype("uint8")
+        masks.update({
+          label : mask
+        })
+      out_hdl = OutputDataHandlerV2()
+      out_hdl.set_masks(w=w, h=h, dict_data=masks, concepts_name2id=self.conceptname2id)
+      out_hdl.set_status(status_code_pb2.SUCCESS)
+      outputs.append(out_hdl._proto)
+
+    return outputs
+
+  def predict(self, request: service_pb2.PostModelOutputsRequest
+             ) -> Iterator[service_pb2.MultiOutputResponse]:
+    """This is the method that will be called when the runner is run. It takes in an input and
+    returns an output.
+    """
+    images = []
+    infer_kwargs = self.get_default_infer_kwargs(request)
+    for input in request.inputs:
+      image = preprocess_image(image_bytes=input.data.image.base64)
+      images.append(image)
+    outputs = self._model_predict(images)
+
+    return service_pb2.MultiOutputResponse(
+        outputs=outputs, status=status_pb2.Status(code=status_code_pb2.SUCCESS)
+        )
+
+  def generate(self, request: service_pb2.PostModelOutputsRequest
+              ) -> Iterator[service_pb2.MultiOutputResponse]:
+    if len(request.inputs) != 1:
+      raise ValueError("Only one input is allowed for image models for this method.")
+    infer_kwargs = self.get_default_infer_kwargs(request)
+
+    for input in request.inputs:
+      input_data = input.data
+      video_bytes = None
+      if input_data.video.base64:
+        video_bytes = input_data.video.base64
+      if video_bytes:
+        frame_generator = video_to_frames(video_bytes)
+        for frame in frame_generator:
+          images = [preprocess_image(frame)]
+          outputs = self._model_predict(images)
+
+        yield service_pb2.MultiOutputResponse(
+            outputs=outputs, status=status_pb2.Status(code=status_code_pb2.SUCCESS))
+
+      else:
+        raise ValueError("Only video input is allowed for this method.")
+
+  def stream(self, request_iterator: Iterator[service_pb2.PostModelOutputsRequest]
+            ) -> Iterator[service_pb2.MultiOutputResponse]:
+    for request in request_iterator:
+      if request.inputs[0].data.video.base64:
+        for output in self.generate(request):
+          yield output
+      elif request.inputs[0].data.image.base64:
+        yield self.predict(request)

models/model_upload/image-segmenter/mask2former-ade/1/utils.py

@@ -0,0 +1,179 @@
+import base64
+from io import BytesIO
+import tempfile
+from typing import Dict, List, Tuple
+
+from google.protobuf import json_format
+import cv2
+from PIL import Image
+import cv2
+import numpy as np
+from clarifai_grpc.grpc.api.status import status_code_pb2
+from clarifai.utils.logging import logger
+from clarifai.runners.utils.data_handler import InputDataHandler, OutputDataHandler
+from clarifai_grpc.grpc.api import resources_pb2
+
+def numpy_image_to_bytes(image_array: np.ndarray, format: str = 'PNG') -> str:
+    if len(image_array.shape) == 2:  # Grayscale image
+        mode = 'L'
+    elif len(image_array.shape) == 3 and image_array.shape[2] == 3:  # RGB image
+        mode = 'RGB'
+    elif len(image_array.shape) == 3 and image_array.shape[2] == 4:  # RGBA image
+        mode = 'RGBA'
+    else:
+        raise ValueError("Unsupported array shape for image conversion.")
+
+    image = Image.fromarray(image_array.astype('uint8'), mode)
+
+    buffer = BytesIO()
+    image.save(buffer, format=format)
+    buffer.seek(0)
+
+    return buffer.getvalue()
+
+def preprocess_image(image_bytes):
+  """Fetch and preprocess image data from bytes"""
+  return Image.open(BytesIO(image_bytes)).convert("RGB")
+
+
+def video_to_frames(video_bytes):
+  """Convert video bytes to frames"""
+  # Write video bytes to a temporary file
+  with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_video_file:
+    temp_video_file.write(video_bytes)
+    temp_video_path = temp_video_file.name
+    logger.info(f"temp_video_path: {temp_video_path}")
+
+    video = cv2.VideoCapture(temp_video_path)
+    print("video opened")
+    logger.info(f"video opened: {video.isOpened()}")
+    while video.isOpened():
+      ret, frame = video.read()
+      if not ret:
+        break
+      # Convert the frame to byte format
+      frame_bytes = cv2.imencode('.jpg', frame)[1].tobytes()
+      yield frame_bytes
+    video.release()
+
+def get_inference_params(request) -> dict:
+  """Get the inference params from the request."""
+  inference_params = {}
+  if request.model.model_version.id != "":
+    output_info = request.model.model_version.output_info
+    output_info = json_format.MessageToDict(output_info, preserving_proto_field_name=True)
+    if "params" in output_info:
+      inference_params = output_info["params"]
+  return inference_params
+
+def create_polygon(mask: np.ndarray) -> List[List[Tuple[float, float]]]:
+  """Create polygons from np mask
+
+  Args:
+      mask (np.ndarray): gray scale binary mask
+
+  Returns:
+      List[List[Tuple[float, float]]]: List of list polygon coordinates
+  """
+  polygons = []
+  contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
+  for obj in contours:
+    coords = []
+    for point in obj:
+      coords.append((int(point[0][0]), int(point[0][1])))
+    polygons.append(coords)
+  return polygons
+
+def make_concept_name(concept):
+  return concept if not concept.startswith("id-") else concept[3:]
+
+class OutputDataHandlerV2(OutputDataHandler):
+  def __init__(self, **kwargs):
+    super().__init__(**kwargs)
+    self.set_status(status_code_pb2.SUCCESS)
+
+  def set_polygons(self, w, h, dict_data: Dict[str, List[np.ndarray]], concepts_name2id: dict = {}, from_type="mask"):
+    types = ["mask", "polygon"]
+    assert from_type in types, ValueError(f"Only accept from_type in {types}. 'mask' is 2d np.array and 'polygon' is 1d array of [x,y,x1,y1..]")
+    regions = []
+    for concept, data in dict_data.items():
+      concept_id = concepts_name2id.get(concept) or make_concept_name(concept)
+      concept_proto = resources_pb2.Concept(
+                    name=concept,
+                    id=concept_id,
+                    value=1.
+                )
+      if from_type == "mask":
+        #mask_h, mask_w = each_data.shape[:2]
+        #assert mask_w == w and mask_h == h, ValueError('Mask size must equal to input image size')
+        polygons = create_polygon(data)
+      elif from_type == "polygon":
+        polygons = data
+      for polygon in polygons:
+        points = []
+        for (x, y) in polygon:
+          x = x / w
+          y = y / h
+          point = resources_pb2.Point(row=x, col=y)
+          points.append(point)
+        region = resources_pb2.Region(
+            region_info=resources_pb2.RegionInfo(polygon=resources_pb2.Polygon(points=points),),
+            data=resources_pb2.Data(concepts=[
+                concept_proto
+            ]))
+        regions.append(region)
+
+    self._proto.data.regions.extend(regions)
+
+  def set_masks(self, w, h, dict_data: Dict[str, List[np.ndarray]], concepts_name2id: dict = {}):
+
+    regions = []
+    for concept, np_mask in dict_data.items():
+      concept_id = concepts_name2id.get(concept) or make_concept_name(concept)
+      concept_proto = resources_pb2.Concept(
+                    name=concept,
+                    id=concept_id,
+                    value=1.
+                )
+      bytes_mask = numpy_image_to_bytes(np_mask, "JPEG")
+      # with open(f".venv/tmp/{concept}.jpg", "wb") as f:
+      #   f.write(bytes_mask)
+      region = resources_pb2.Region(
+          region_info = resources_pb2.RegionInfo(
+            mask = resources_pb2.Mask(
+              image = resources_pb2.Image(base64=bytes_mask)
+            )
+          ),
+          data = resources_pb2.Data(concepts=[
+              concept_proto
+          ]))
+      regions.append(region)
+
+    self._proto.data.regions.extend(regions)
+
+
+
+
+###############
+import yaml
+from transformers import AutoConfig
+
+def create_concepts_in_yaml(path, model_short_name, repo_id=None):
+  with open(path, "r") as f:
+    data = yaml.safe_load(f)
+
+  if not repo_id:
+    repo_id = data["checkpoints"]["repo_id"]
+  model_config = AutoConfig.from_pretrained(repo_id)
+  concepts = []
+  for _id, lb in model_config.id2label.items():
+    concepts.append({
+      "id": f"id-{model_short_name}-{_id}",
+      "name": lb
+    })
+
+  data.update(dict(concepts=concepts))
+  with open(path, "w") as f:
+    yaml.dump(data, f, sort_keys=False)
+
+  return data