File size: 4,913 Bytes
62fd502 |
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 142 143 144 145 146 147 148 149 150 151 152 153 |
import cv2
import numpy as np
from PIL import Image
import streamlit as st
import tensorflow as tf
from tensorflow.keras.models import load_model
# most of this code has been obtained from Datature's prediction script
# https://github.com/datature/resources/blob/main/scripts/bounding_box/prediction.py
st.set_option('deprecation.showfileUploaderEncoding', False)
@st.cache(allow_output_mutation=True)
def load_model():
return tf.saved_model.load('./saved_model')
def load_label_map(label_map_path):
"""
Reads label map in the format of .pbtxt and parse into dictionary
Args:
label_map_path: the file path to the label_map
Returns:
dictionary with the format of {label_index: {'id': label_index, 'name': label_name}}
"""
label_map = {}
with open(label_map_path, "r") as label_file:
for line in label_file:
if "id" in line:
label_index = int(line.split(":")[-1])
label_name = next(label_file).split(":")[-1].strip().strip('"')
label_map[label_index] = {"id": label_index, "name": label_name}
return label_map
def predict_class(image, model):
image = tf.cast(image, tf.float32)
image = tf.image.resize(image, [150, 150])
image = np.expand_dims(image, axis = 0)
return model.predict(image)
def plot_boxes_on_img(color_map, classes, bboxes, image_origi, origi_shape):
for idx, each_bbox in enumerate(bboxes):
color = color_map[classes[idx]]
## Draw bounding box
cv2.rectangle(
image_origi,
(int(each_bbox[1] * origi_shape[1]),
int(each_bbox[0] * origi_shape[0]),),
(int(each_bbox[3] * origi_shape[1]),
int(each_bbox[2] * origi_shape[0]),),
color,
2,
)
## Draw label background
cv2.rectangle(
image_origi,
(int(each_bbox[1] * origi_shape[1]),
int(each_bbox[2] * origi_shape[0]),),
(int(each_bbox[3] * origi_shape[1]),
int(each_bbox[2] * origi_shape[0] + 15),),
color,
-1,
)
## Insert label class & score
cv2.putText(
image_origi,
"Class: {}, Score: {}".format(
str(category_index[classes[idx]]["name"]),
str(round(scores[idx], 2)),
),
(int(each_bbox[1] * origi_shape[1]),
int(each_bbox[2] * origi_shape[0] + 10),),
cv2.FONT_HERSHEY_SIMPLEX,
0.3,
(0, 0, 0),
1,
cv2.LINE_AA,
)
return image_origi
# Webpage code starts here
#TODO change this
st.title('YOUR PROJECT NAME')
st.text('made by XXX')
st.markdown('## Description about your project')
with st.spinner('Model is being loaded...'):
model = load_model()
# ask user to upload an image
file = st.file_uploader("Upload image", type=["jpg", "png"])
if file is None:
st.text('Waiting for upload...')
else:
st.text('Running inference...')
# open image
test_image = Image.open(file).convert("RGB")
origi_shape = np.asarray(test_image).shape
# resize image to default shape
default_shape = 320
image_resized = np.array(test_image.resize((default_shape, default_shape)))
## Load color map
category_index = load_label_map("./label_map.pbtxt")
# TODO Add more colors if there are more classes
# color of each label. check label_map.pbtxt to check the index for each class
color_map = {
1: [255, 0, 0], # bad -> red
2: [0, 255, 0] # good -> green
}
## The model input needs to be a tensor
input_tensor = tf.convert_to_tensor(image_resized)
## The model expects a batch of images, so add an axis with `tf.newaxis`.
input_tensor = input_tensor[tf.newaxis, ...]
## Feed image into model and obtain output
detections_output = model(input_tensor)
num_detections = int(detections_output.pop("num_detections"))
detections = {key: value[0, :num_detections].numpy() for key, value in detections_output.items()}
detections["num_detections"] = num_detections
## Filter out predictions below threshold
# if threshold is higher, there will be fewer predictions
# TODO change this number to see how the predictions change
confidence_threshold = 0.8
indexes = np.where(detections["detection_scores"] > confidence_threshold)
## Extract predicted bounding boxes
bboxes = detections["detection_boxes"][indexes]
# there are no predicted boxes
if len(bboxes) == 0:
st.error('No boxes predicted')
# there are predicted boxes
else:
st.success('Boxes predicted')
classes = detections["detection_classes"][indexes].astype(np.int64)
scores = detections["detection_scores"][indexes]
# plot boxes and labels on image
image_origi = np.array(Image.fromarray(image_resized).resize((origi_shape[1], origi_shape[0])))
image_origi = plot_boxes_on_img(color_map, classes, bboxes, image_origi, origi_shape)
# show image in web page
st.image(Image.fromarray(image_origi), caption="Image with predictions", width=400)
st.markdown("### Predicted boxes")
for idx in range(len((bboxes))):
st.markdown(f"* Class: {str(category_index[classes[idx]]['name'])}, confidence score: {str(round(scores[idx], 2))}") |