| import json | |
| from typing import Any, Dict, List | |
| import tensorflow as tf | |
| from tensorflow import keras | |
| import base64 | |
| import io | |
| import os | |
| import numpy as np | |
| from PIL import Image | |
| # most of this code has been obtained from Datature's prediction script | |
| # https://github.com/datature/resources/blob/main/scripts/bounding_box/prediction.py | |
| # def load_model(): | |
| # return tf.saved_model.load('./saved_model') | |
| # model = load_model() | |
| class PreTrainedPipeline(): | |
| def __init__(self, path: str): | |
| # load the model | |
| self.model = keras.models.load_model(os.path.join(path, "tf_model.h5")) | |
| #self.model = tf.saved_model.load(os.path.join(path, "saved_model")) | |
| def __call__(self, inputs: "Image.Image")-> List[Dict[str, Any]]: | |
| # # convert img to numpy array, resize and normalize to make the prediction | |
| # img = np.array(inputs) | |
| # im = tf.image.resize(img, (128, 128)) | |
| # im = tf.cast(im, tf.float32) / 255.0 | |
| # pred_mask = self.model.predict(im[tf.newaxis, ...]) | |
| # # take the best performing class for each pixel | |
| # # the output of argmax looks like this [[1, 2, 0], ...] | |
| # pred_mask_arg = tf.argmax(pred_mask, axis=-1) | |
| # labels = [] | |
| # # convert the prediction mask into binary masks for each class | |
| # binary_masks = {} | |
| # mask_codes = {} | |
| # # when we take tf.argmax() over pred_mask, it becomes a tensor object | |
| # # the shape becomes TensorShape object, looking like this TensorShape([128]) | |
| # # we need to take get shape, convert to list and take the best one | |
| # rows = pred_mask_arg[0][1].get_shape().as_list()[0] | |
| # cols = pred_mask_arg[0][2].get_shape().as_list()[0] | |
| # for cls in range(pred_mask.shape[-1]): | |
| # binary_masks[f"mask_{cls}"] = np.zeros(shape = (pred_mask.shape[1], pred_mask.shape[2])) #create masks for each class | |
| # for row in range(rows): | |
| # for col in range(cols): | |
| # if pred_mask_arg[0][row][col] == cls: | |
| # binary_masks[f"mask_{cls}"][row][col] = 1 | |
| # else: | |
| # binary_masks[f"mask_{cls}"][row][col] = 0 | |
| # mask = binary_masks[f"mask_{cls}"] | |
| # mask *= 255 | |
| # img = Image.fromarray(mask.astype(np.int8), mode="L") | |
| # # we need to make it readable for the widget | |
| # with io.BytesIO() as out: | |
| # img.save(out, format="PNG") | |
| # png_string = out.getvalue() | |
| # mask = base64.b64encode(png_string).decode("utf-8") | |
| # mask_codes[f"mask_{cls}"] = mask | |
| # # widget needs the below format, for each class we return label and mask string | |
| # labels.append({ | |
| # "label": f"LABEL_{cls}", | |
| # "mask": mask_codes[f"mask_{cls}"], | |
| # "score": 1.0, | |
| # }) | |
| # labels = [{"score":0.9509243965148926,"label":"car","box":{"xmin":142,"ymin":106,"xmax":376,"ymax":229}}, | |
| # {"score":0.9981777667999268,"label":"car","box":{"xmin":405,"ymin":146,"xmax":640,"ymax":297}}, | |
| # {"score":0.9963648915290833,"label":"car","box":{"xmin":0,"ymin":115,"xmax":61,"ymax":167}}, | |
| # {"score":0.974663257598877,"label":"car","box":{"xmin":155,"ymin":104,"xmax":290,"ymax":141}}, | |
| # {"score":0.9986898303031921,"label":"car","box":{"xmin":39,"ymin":117,"xmax":169,"ymax":188}}, | |
| # {"score":0.9998276233673096,"label":"person","box":{"xmin":172,"ymin":60,"xmax":482,"ymax":396}}, | |
| # {"score":0.9996274709701538,"label":"skateboard","box":{"xmin":265,"ymin":348,"xmax":440,"ymax":413}}] | |
| labels = [] | |
| return labels | |