File size: 2,218 Bytes
9f07f93 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import base64
import torch
import numpy as np
from PIL import Image
import io
class BaseHandler:
def __init__(self):
"""Initialize the handler with model-specific configurations"""
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = None
self.initialized = False
def initialize(self):
"""Load model and other resources"""
# This method should be implemented by each specific handler
raise NotImplementedError
def preprocess(self, data):
"""Preprocess the input data"""
# This method should be implemented by each specific handler
raise NotImplementedError
def inference(self, inputs):
"""Run inference with the preprocessed inputs"""
# This method should be implemented by each specific handler
raise NotImplementedError
def postprocess(self, inference_output):
"""Post-process the model output"""
# This method should be implemented by each specific handler
raise NotImplementedError
def __call__(self, data):
"""Handle a request to the model"""
# Initialize the model if not already done
if not self.initialized:
self.initialize()
# Process the request
preprocessed_data = self.preprocess(data)
inference_output = self.inference(preprocessed_data)
output = self.postprocess(inference_output)
return output
def encode_image(self, image):
"""Encode a PIL Image to base64"""
buffered = io.BytesIO()
image.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
return img_str
def decode_image(self, image_str):
"""Decode a base64 string to PIL Image"""
img_data = base64.b64decode(image_str)
return Image.open(io.BytesIO(img_data))
def svg_to_base64(self, svg_content):
"""Convert SVG content to base64"""
return base64.b64encode(svg_content.encode("utf-8")).decode("utf-8") |