Upload 15 files

0ec106b verified about 2 years ago

2.96 kB

import os
import numpy as np
from typing import Dict

from transformers import MaskFormerFeatureExtractor
# from refactor.utils.assets import ImageAsset
from io import BytesIO
from PIL import Image
import rawpy

import triton_python_backend_utils as pb_utils


class TritonPythonModel:
    def initialize(self, args: Dict[str, str]) -> None:
        """
        Initialize the tokenization process
        :param args: arguments from Triton config file
        """
        path: str = os.path.join(args["model_repository"], args["model_version"])
        self.processor = MaskFormerFeatureExtractor.from_pretrained(path)
        
    def execute(self, requests):
        """
        Parameters
        ----------
        requests : list
          A list of pb_utils.InferenceRequest
        Returns
        -------
        list
          A list of pb_utils.InferenceResponse. The length of this list must
          be the same as `requests`
        """
        responses = []
        # loop batch request
        for request in requests:
            # Get INPUT
            # img = [
            #     self.preprocess(t)
            #     for t in pb_utils.get_input_tensor_by_name(request, "IMAGE")
            #     .as_numpy()
            # ]
            imgs = pb_utils.get_input_tensor_by_name(request, "IMAGE").as_numpy().tolist()
            imgs = self.preprocess(imgs)
            # imgs = [Image.open(BytesIO(img)) for img in imgs]
            # img = ImageAsset(inp).image
            # img = [self.preprocess(e) for e in inp]

            # Get EXTRACTFEATURE
            feature: Dict[str, np.ndarray] = self.processor(imgs, return_tensors='np')
            outputs = pb_utils.Tensor('pixel_values', feature['pixel_values'])
            inference_response = pb_utils.InferenceResponse(output_tensors=[outputs])
            responses.append(inference_response)
            
        return responses
    

    def preprocess(self, data):
      print("images: ", data)
      imgs = []
      for e in data:
        buffer = BytesIO(e)
        print("Type of data: ", type(buffer))
        _image = Image.open(buffer).convert('RGB')
        print("Type of _image: ", type(_image))
        imgs.append(_image)
      print("Load success")
      # try:
      #   with rawpy.imread(buffer) as raw:
      #     thumb = raw.extract_thumb()
      #     print("Load success: ",thumb)
      #   if thumb.format == rawpy.ThumbFormat.JPEG or thumb.format == rawpy.ThumbFormat.PNG:
      #     _image = Image.open(thumb.data)
      #     print("PNG Type of _image: ", type(_image))
      #   elif thumb.format == rawpy.ThumbFormat.BITMAP:
      #     _image = Image.fromarray(thumb.data)
      #     print("bitmap Type of _image: ", type(_image))
      # except:
      #   raise AssertionError(f'Cannot load image asset from {str(buffer):20s}')
      return imgs
      
    def finalize(self) -> None:
        """
        Finalize the model
        """
        pass

	import os
	import numpy as np
	from typing import Dict

	from transformers import MaskFormerFeatureExtractor
	# from refactor.utils.assets import ImageAsset
	from io import BytesIO
	from PIL import Image
	import rawpy

	import triton_python_backend_utils as pb_utils


	class TritonPythonModel:
	def initialize(self, args: Dict[str, str]) -> None:
	"""
	Initialize the tokenization process
	:param args: arguments from Triton config file
	"""
	path: str = os.path.join(args["model_repository"], args["model_version"])
	self.processor = MaskFormerFeatureExtractor.from_pretrained(path)

	def execute(self, requests):
	"""
	Parameters
	----------
	requests : list
	A list of pb_utils.InferenceRequest

	Returns
	-------
	list
	A list of pb_utils.InferenceResponse. The length of this list must
	be the same as `requests`
	"""
	responses = []
	# loop batch request
	for request in requests:
	# Get INPUT
	# img = [
	# self.preprocess(t)
	# for t in pb_utils.get_input_tensor_by_name(request, "IMAGE")
	# .as_numpy()
	# ]
	imgs = pb_utils.get_input_tensor_by_name(request, "IMAGE").as_numpy().tolist()
	imgs = self.preprocess(imgs)
	# imgs = [Image.open(BytesIO(img)) for img in imgs]
	# img = ImageAsset(inp).image
	# img = [self.preprocess(e) for e in inp]

	# Get EXTRACTFEATURE
	feature: Dict[str, np.ndarray] = self.processor(imgs, return_tensors='np')
	outputs = pb_utils.Tensor('pixel_values', feature['pixel_values'])
	inference_response = pb_utils.InferenceResponse(output_tensors=[outputs])
	responses.append(inference_response)

	return responses


	def preprocess(self, data):
	print("images: ", data)
	imgs = []
	for e in data:
	buffer = BytesIO(e)
	print("Type of data: ", type(buffer))
	_image = Image.open(buffer).convert('RGB')
	print("Type of _image: ", type(_image))
	imgs.append(_image)
	print("Load success")
	# try:
	# with rawpy.imread(buffer) as raw:
	# thumb = raw.extract_thumb()
	# print("Load success: ",thumb)
	# if thumb.format == rawpy.ThumbFormat.JPEG or thumb.format == rawpy.ThumbFormat.PNG:
	# _image = Image.open(thumb.data)
	# print("PNG Type of _image: ", type(_image))
	# elif thumb.format == rawpy.ThumbFormat.BITMAP:
	# _image = Image.fromarray(thumb.data)
	# print("bitmap Type of _image: ", type(_image))
	# except:
	# raise AssertionError(f'Cannot load image asset from {str(buffer):20s}')
	return imgs

	def finalize(self) -> None:
	"""
	Finalize the model
	"""
	pass