Spaces:
Runtime error
Runtime error
File size: 5,099 Bytes
db5855f |
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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
# SegmentationModel implementation based on
# https://github.com/openvinotoolkit/open_model_zoo/tree/master/demos/common/python/models
import sys
import cv2
import numpy as np
from os import PathLike
from openvino.runtime import PartialShape
import logging
# Fetch `notebook_utils` module
import requests
r = requests.get(
url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py",
)
open("notebook_utils.py", "w").write(r.text)
from notebook_utils import segmentation_map_to_overlay
def sigmoid(x):
return np.exp(-np.logaddexp(0, -x))
class Model:
def __init__(self, ie, model_path, input_transform=None):
self.logger = logging.getLogger()
self.logger.info("Reading model from IR...")
self.net = ie.read_model(model_path)
self.set_batch_size(1)
self.input_transform = input_transform
def preprocess(self, inputs):
meta = {}
return inputs, meta
def postprocess(self, outputs, meta):
return outputs
def set_batch_size(self, batch):
shapes = {}
for input_layer in self.net.inputs:
new_shape = input_layer.partial_shape
new_shape[0] = 1
shapes.update({input_layer: new_shape})
self.net.reshape(shapes)
class SegmentationModel(Model):
def __init__(
self,
ie,
model_path: PathLike,
colormap: np.ndarray = None,
resize_shape=None,
sigmoid=False,
argmax=False,
rgb=False,
rotate_and_flip=False,
):
"""
Segmentation Model for use with Async Pipeline.
:param model_path: path to IR model .xml file
:param colormap: array of shape (num_classes, 3) where colormap[i] contains the RGB color
values for class i. Optional for binary segmentation, required for multiclass
:param resize_shape: if specified, reshape the model to this shape
:param sigmoid: if True, apply sigmoid to model result
:param argmax: if True, apply argmax to model result
:param rgb: set to True if the model expects RGB images as input
"""
super().__init__(ie, model_path)
self.sigmoid = sigmoid
self.argmax = argmax
self.rgb = rgb
self.rotate_and_flip = rotate_and_flip
self.net = ie.read_model(model_path)
self.input_layer = self.net.input(0)
self.output_layer = self.net.output(0)
if resize_shape is not None:
self.net.reshape({self.input_layer: PartialShape(resize_shape)})
self.image_height, self.image_width = (
self.input_layer.shape[2],
self.input_layer.shape[3],
)
if colormap is None and self.output_layer.shape[1] == 1:
self.colormap = np.array([[0, 0, 0], [0, 0, 255]])
else:
self.colormap = colormap
if self.colormap is None:
raise ValueError("Please provide a colormap for multiclass segmentation")
def preprocess(self, inputs):
"""
Resize the image to network input dimensions and transpose to
network input shape with N,C,H,W layout.
Images are expected to have dtype np.uint8 and shape (H,W,3) or (H,W)
"""
meta = {}
image = inputs[self.input_layer]
meta["frame"] = image
if image.shape[:2] != (self.image_height, self.image_width):
image = cv2.resize(image, (self.image_width, self.image_height))
if len(image.shape) == 3:
input_image = np.expand_dims(np.transpose(image, (2, 0, 1)), 0)
else:
input_image = np.expand_dims(np.expand_dims(image, 0), 0)
return {0: input_image}, meta
def postprocess(self, outputs, preprocess_meta, to_rgb=False):
"""
Convert raw network results into a segmentation map with overlay. Returns
a BGR image for further processing with OpenCV.
"""
alpha = 0.4
if preprocess_meta["frame"].shape[-1] == 3:
bgr_frame = preprocess_meta["frame"]
if self.rgb:
# reverse color channels to convert to BGR
bgr_frame = bgr_frame[:, :, (2, 1, 0)]
else:
# Create BGR image by repeating channels in one-channel image
bgr_frame = np.repeat(np.expand_dims(preprocess_meta["frame"], -1), 3, 2)
res = outputs[0].squeeze()
result_mask_ir = sigmoid(res) if self.sigmoid else res
if self.argmax:
result_mask_ir = np.argmax(res, axis=0).astype(np.uint8)
else:
result_mask_ir = result_mask_ir.round().astype(np.uint8)
overlay = segmentation_map_to_overlay(bgr_frame, result_mask_ir, alpha, colormap=self.colormap)
if self.rotate_and_flip:
overlay = cv2.flip(cv2.rotate(overlay, rotateCode=cv2.ROTATE_90_CLOCKWISE), flipCode=1)
return overlay
|