app.py

app.py

@@ -0,0 +1,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))}")