|
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 |
|
|
|
|
|
|
|
|
|
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]] |
|
|
|
|
|
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, |
|
) |
|
|
|
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, |
|
) |
|
|
|
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 |
|
|
|
|
|
|
|
|
|
|
|
st.title('Distribution Grid - Belgium - Equipment detection') |
|
st.text('made by LabelFlow') |
|
st.markdown('## Description about your project') |
|
|
|
with st.spinner('Model is being loaded...'): |
|
model = load_model() |
|
|
|
|
|
file = st.file_uploader("Upload image", type=["jpg", "png"]) |
|
|
|
if file is None: |
|
st.text('Waiting for upload...') |
|
else: |
|
st.text('Running inference...') |
|
|
|
test_image = Image.open(file).convert("RGB") |
|
origi_shape = np.asarray(test_image).shape |
|
|
|
default_shape = 320 |
|
image_resized = np.array(test_image.resize((default_shape, default_shape))) |
|
|
|
|
|
category_index = load_label_map("./label_map.pbtxt") |
|
|
|
|
|
|
|
color_map = { |
|
1: [69, 109, 42], |
|
2: [107, 46, 186], |
|
3: [9, 35, 183], |
|
4: [27, 1, 30], |
|
5: [0, 0, 0], |
|
6: [5, 6, 7], |
|
7: [11, 5, 12], |
|
8: [209, 205, 211], |
|
9: [17, 17, 17], |
|
10: [101, 242, 50], |
|
11: [51, 204, 170], |
|
12: [106, 0, 132], |
|
13: [7, 111, 153], |
|
14: [8, 10, 9], |
|
15: [234, 250, 252], |
|
16: [58, 68, 30], |
|
17: [24, 178, 117], |
|
18: [21, 22, 21], |
|
19: [53, 104, 83], |
|
20: [12, 5, 10], |
|
21: [223, 192, 249], |
|
22: [234, 234, 234], |
|
23: [119, 68, 221], |
|
24: [224, 174, 94], |
|
25: [140, 74, 116], |
|
26: [90, 102, 1], |
|
27: [216, 143, 208] |
|
} |
|
|
|
|
|
input_tensor = tf.convert_to_tensor(image_resized) |
|
|
|
input_tensor = input_tensor[tf.newaxis, ...] |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
|
confidence_threshold = 0.6 |
|
indexes = np.where(detections["detection_scores"] > confidence_threshold) |
|
|
|
|
|
bboxes = detections["detection_boxes"][indexes] |
|
|
|
if len(bboxes) == 0: |
|
st.error('No boxes predicted') |
|
|
|
else: |
|
st.success('Boxes predicted') |
|
classes = detections["detection_classes"][indexes].astype(np.int64) |
|
scores = detections["detection_scores"][indexes] |
|
|
|
|
|
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) |
|
|
|
|
|
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))}") |
|
|