pytorch / serve

Serve, optimize and scale PyTorch models in production
https://pytorch.org/serve/
Apache License 2.0
4.17k stars 849 forks source link

Custom Handler Detectron2 #738

Closed reubenwenisch closed 3 years ago

reubenwenisch commented 3 years ago

Context

I am trying to serve a detectron2 model with torch serve using a custom handler. I get the following error "TypeError: Object of type Instances is not JSON serializable" after I convert the dictionary output of the model to a list.

Custom Handler Code

from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2 import model_zoo
import torch, os

from ts.torch_handler.base_handler import BaseHandler
from torchvision import transforms as T
from PIL import Image
import io
import numpy

cfg = get_cfg()

cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.1   # set the testing threshold for this model
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 7  # only has one class (ballon)

class ModelHandler(object):
    """
    A custom model handler implementation.
    """
    image_processing = T.Compose([
        T.Resize(256),
        T.CenterCrop(224),
        T.ToTensor(),
        T.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )])

    def __init__(self):
        self._context = None
        self.initialized = False
        self.model = None
        self.device = None

    def initialize(self, context):
        """
        Invoke by torchserve for loading a model
        :param context: context contains model server system properties
        :return:
        """

        #  load the model
        self.manifest = context.manifest

        properties = context.system_properties
        model_dir = properties.get("model_dir")
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        print(" device is selected")
        # Read model serialize/pt file
        serialized_file = self.manifest['model']['serializedFile']
        model_pt_path = os.path.join(model_dir, serialized_file)
        if not os.path.isfile(model_pt_path):
            raise RuntimeError("Missing the model.pt file")
        cfg.MODEL.WEIGHTS = model_pt_path
        self.model = DefaultPredictor(cfg)

        self.initialized = True

    def preprocess(self, data):
        for row in data:
            image = row.get("data") or row.get("body")
            image = Image.open(io.BytesIO(image))
            open_cv_image = numpy.array(image) 
            # Convert RGB to BGR 
            open_cv_image = open_cv_image[:, :, ::-1].copy() 
        return open_cv_image

    def postprocess(self, data):
        result = list(data.items())
        return result

    def handle(self, data, context):
        """
        Invoke by TorchServe for prediction request.
        Do pre-processing of data, prediction using model and postprocessing of prediciton output
        :param data: Input data for prediction
        :param context: Initial context contains model server system properties.
        :return: prediction output
        """
        model_input = self.preprocess(data)
        model_output = self.model(model_input)
        return self.postprocess(model_output)`

Expected Behavior

Want the dictionary contents via api.

Current Behavior

raise TypeError(f'Object of type {o.class.name} ' TypeError: Object of type Instances is not JSON serializable

harshbafna commented 3 years ago

@reubenwenisch: The DefaultPredictor model returns object of type Instances and not a JSON/python dictionary.

Want the dictionary contents via api.

I assume you are looking for boxes and classes. You can extract these from the model output as follows :

outputs = model(image)
print(outputs['instances'].pred_boxes)
print(outputs['instances'].pred_classes)
reubenwenisch commented 3 years ago

Right. Even this pred_boxes is of type detectron2.structures.boxes. Even this cannot be json serializable right? So what is the recommended way to send the output and convert back in the case of a segmentation algorithm? Currently, I am converting to a list is that the recommended way?

harshbafna commented 3 years ago

@reubenwenisch: As documented by detectron2, detectron2.structures.boxes stores a list of boxes as an Nx4 torch.Tensor.

So what is the recommended way to send the output and convert back in the case of a segmentation algorithm TorchServe handler should return a list of output. You can return any type of data like Tensor, Boxes, Classes, or a segmented image (in binary format). What you want to return from the handler will depend on your use-case at the client end.

You can refer to the following google-colab jupyter notebook which demonstrates running inference on the DefaultPredictor

https://colab.research.google.com/drive/16jcaJoc6bCFAQ96jDe2HwtXj7BMD_-m5#scrollTo=7d3KxiHO_0gb

You can also refer to the post-processing of TorchServe's default image-segmenter handler