Update pipeline.py
Browse files- pipeline.py +22 -10
pipeline.py
CHANGED
@@ -13,31 +13,43 @@ from PIL import Image
|
|
13 |
|
14 |
class PreTrainedPipeline():
|
15 |
def __init__(self, path: str):
|
16 |
-
|
17 |
self.model = keras.models.load_model(os.path.join(path, "tf_model.h5"))
|
18 |
|
19 |
def __call__(self, inputs: "Image.Image")-> List[Dict[str, Any]]:
|
20 |
|
|
|
21 |
with Image.open(inputs) as img:
|
22 |
img = np.array(img)
|
23 |
|
24 |
im = tf.image.resize(img, (128, 128))
|
25 |
im = tf.cast(im, tf.float32) / 255.0
|
26 |
pred_mask = model.predict(im[tf.newaxis, ...])
|
|
|
|
|
|
|
27 |
pred_mask_arg = tf.argmax(pred_mask, axis=-1)
|
28 |
|
29 |
labels = []
|
30 |
-
|
|
|
31 |
binary_masks = {}
|
32 |
mask_codes = {}
|
33 |
-
|
34 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
35 |
for cls in range(pred_mask.shape[-1]):
|
36 |
|
37 |
-
binary_masks[f"mask_{cls}"] = np.zeros(shape = (pred_mask.shape[1], pred_mask.shape[2]))
|
38 |
-
|
|
|
39 |
|
40 |
-
for col in range(
|
41 |
|
42 |
if pred_mask_arg[0][row][col] == cls:
|
43 |
|
@@ -48,7 +60,8 @@ class PreTrainedPipeline():
|
|
48 |
mask = binary_masks[f"mask_{cls}"]
|
49 |
mask *= 255
|
50 |
img = Image.fromarray(mask.astype(np.int8), mode="L")
|
51 |
-
|
|
|
52 |
with io.BytesIO() as out:
|
53 |
img.save(out, format="PNG")
|
54 |
png_string = out.getvalue()
|
@@ -56,9 +69,8 @@ class PreTrainedPipeline():
|
|
56 |
|
57 |
mask_codes[f"mask_{cls}"] = mask
|
58 |
|
59 |
-
|
60 |
|
61 |
-
|
62 |
labels.append({
|
63 |
"label": f"LABEL_{cls}",
|
64 |
"mask": mask_codes[f"mask_{cls}"],
|
|
|
13 |
|
14 |
class PreTrainedPipeline():
|
15 |
def __init__(self, path: str):
|
16 |
+
# load the model
|
17 |
self.model = keras.models.load_model(os.path.join(path, "tf_model.h5"))
|
18 |
|
19 |
def __call__(self, inputs: "Image.Image")-> List[Dict[str, Any]]:
|
20 |
|
21 |
+
# convert img to numpy array, resize and normalize to make the prediction
|
22 |
with Image.open(inputs) as img:
|
23 |
img = np.array(img)
|
24 |
|
25 |
im = tf.image.resize(img, (128, 128))
|
26 |
im = tf.cast(im, tf.float32) / 255.0
|
27 |
pred_mask = model.predict(im[tf.newaxis, ...])
|
28 |
+
|
29 |
+
# take the best performing class for each pixel
|
30 |
+
# the output of argmax looks like this [[1, 2, 0], ...]
|
31 |
pred_mask_arg = tf.argmax(pred_mask, axis=-1)
|
32 |
|
33 |
labels = []
|
34 |
+
|
35 |
+
# convert the prediction mask into binary masks for each class
|
36 |
binary_masks = {}
|
37 |
mask_codes = {}
|
38 |
+
|
39 |
+
# when we take tf.argmax() over pred_mask, it becomes a tensor object
|
40 |
+
# the shape becomes TensorShape object, looking like this TensorShape([128])
|
41 |
+
# we need to take get shape, convert to list and take the best one
|
42 |
+
|
43 |
+
rows = pred_mask_arg[0][1].get_shape().as_list()[0]
|
44 |
+
cols = pred_mask_arg[0][2].get_shape().as_list()[0]
|
45 |
+
|
46 |
for cls in range(pred_mask.shape[-1]):
|
47 |
|
48 |
+
binary_masks[f"mask_{cls}"] = np.zeros(shape = (pred_mask.shape[1], pred_mask.shape[2])) #create masks for each class
|
49 |
+
|
50 |
+
for row in range(rows):
|
51 |
|
52 |
+
for col in range(cols):
|
53 |
|
54 |
if pred_mask_arg[0][row][col] == cls:
|
55 |
|
|
|
60 |
mask = binary_masks[f"mask_{cls}"]
|
61 |
mask *= 255
|
62 |
img = Image.fromarray(mask.astype(np.int8), mode="L")
|
63 |
+
|
64 |
+
# we need to make it readable for the widget
|
65 |
with io.BytesIO() as out:
|
66 |
img.save(out, format="PNG")
|
67 |
png_string = out.getvalue()
|
|
|
69 |
|
70 |
mask_codes[f"mask_{cls}"] = mask
|
71 |
|
|
|
72 |
|
73 |
+
# widget needs the below format, for each class we return label and mask string
|
74 |
labels.append({
|
75 |
"label": f"LABEL_{cls}",
|
76 |
"mask": mask_codes[f"mask_{cls}"],
|