Spaces:
Sleeping
Sleeping
File size: 3,236 Bytes
4c38319 |
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 |
import numpy as np
import gradio as gr
from tensorflow.keras.models import load_model
import imutils
import matplotlib.pyplot as plt
import cv2
import numpy as np
from tensorflow.keras.preprocessing.image import img_to_array
from PIL import Image
PROB_THRESHOLD = 0.4 # Minimum probably to show results.
class Model:
def __init__(self, model_filepath):
self.graph_def = tensorflow.compat.v1.GraphDef()
self.graph_def.ParseFromString(model_filepath.read_bytes())
input_names, self.output_names = self._get_graph_inout(self.graph_def)
assert len(input_names) == 1 and len(self.output_names) == 3
self.input_name = input_names[0]
self.input_shape = self._get_input_shape(self.graph_def, self.input_name)
def predict(self, image_filepath):
image = Image.fromarray(image_filepath).resize(self.input_shape)
input_array = np.array(image, dtype=np.float32)[np.newaxis, :, :, :]
with tensorflow.compat.v1.Session() as sess:
tensorflow.import_graph_def(self.graph_def, name='')
out_tensors = [sess.graph.get_tensor_by_name(o + ':0') for o in self.output_names]
outputs = sess.run(out_tensors, {self.input_name + ':0': input_array})
return {name: outputs[i][np.newaxis, ...] for i, name in enumerate(self.output_names)}
@staticmethod
def _get_graph_inout(graph_def):
input_names = []
inputs_set = set()
outputs_set = set()
for node in graph_def.node:
if node.op == 'Placeholder':
input_names.append(node.name)
for i in node.input:
inputs_set.add(i.split(':')[0])
outputs_set.add(node.name)
output_names = list(outputs_set - inputs_set)
return input_names, output_names
@staticmethod
def _get_input_shape(graph_def, input_name):
for node in graph_def.node:
if node.name == input_name:
return [dim.size for dim in node.attr['shape'].shape.dim][1:3]
def print_outputs(outputs, gambar):
image = gambar
assert set(outputs.keys()) == set(['detected_boxes', 'detected_classes', 'detected_scores'])
l, t, d = image.shape
labelopen = open("labels.txt", 'r')
labels = [line.split(',') for line in labelopen.readlines()]
for box, class_id, score in zip(outputs['detected_boxes'][0], outputs['detected_classes'][0], outputs['detected_scores'][0]):
if score > PROB_THRESHOLD:
print(f"Label: {class_id}, Probability: {score:.5f}, box: ({box[0]:.5f}, {box[1]:.5f}) ({box[2]:.5f}, {box[3]:.5f})")
x = box[0] * t
y = box[1] * l
h = box[2] * t
w = box[3] * l
result_image = cv2.rectangle(image, (int(x), int(y)), (int(h), int(w)), (255,215,0), 3)
cv2.putText(result_image, labels[int(class_id)][0], (int(x), int(y)-10), fontFace = cv2.FONT_HERSHEY_SIMPLEX, fontScale = 0.5, color = (255,215,0), thickness = 2)
return result_image
def main(gambar):
p = pathlib.Path("model.pb")
#b = pathlib.Path(gambar)
model = Model(p)
outputs = model.predict(gambar)
return print_outputs(outputs, gambar)
demo = gr.Interface(main, gr.Image(shape=(500, 500)), "image")
demo.launch() |