Created
July 8, 2020 12:13
-
-
Save binarymason/288b9bfe0633efad6bf95e3ef49d43d9 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import numpy as np | |
import os | |
import sys | |
import tensorflow as tf | |
import glob | |
from PIL import Image | |
from tensorflow_utils import label_map_util, ops as utils_ops | |
import cv2 | |
# from object_detection.utils import ops as utils_ops | |
MODEL_NAME = 'solar_panel_inference_graph' | |
# Path to frozen detection graph. This is the actual model that is used for the object detection. | |
PATH_TO_FROZEN_GRAPH = MODEL_NAME + '/frozen_inference_graph.pb' | |
# List of the strings that is used to add correct label for each box. | |
# PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt') | |
PATH_TO_LABELS = os.path.join('solar_panel_training', 'labelmap.pbtxt') | |
DETECTION_GRAPH = tf.Graph() | |
with DETECTION_GRAPH.as_default(): | |
od_graph_def = tf.GraphDef() | |
with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid: | |
serialized_graph = fid.read() | |
od_graph_def.ParseFromString(serialized_graph) | |
tf.import_graph_def(od_graph_def, name='') | |
category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True) | |
def load_image_into_numpy_array(image): | |
(im_width, im_height) = image.size | |
return np.array(image.getdata()).reshape( | |
(im_height, im_width, 3)).astype(np.uint8) | |
PATH_TO_TEST_IMAGES_DIR = "/home/m/code/banshee/test_images/videos/frames" | |
TEST_IMAGE_PATHS = [file for file in glob.glob(PATH_TO_TEST_IMAGES_DIR + "/*.jpg")] | |
def detect(path, score_threshold = 0.5): | |
img = cv2.imread(path) | |
out = run_inference_for_single_image(path) | |
rows, cols, _ = img.shape | |
result = { "solar_panel": [], "hotspot": [] } | |
labelmap = { 1: "solar_panel", 2: "hotspot" } | |
num_detections = int(out["num_detections"]) | |
for i in range(num_detections): | |
classId = int(out["detection_classes"][i]) | |
label = labelmap[classId] | |
score = float(out["detection_scores"][i]) | |
bbox = [float(v) for v in out["detection_boxes"][i]] | |
if score > score_threshold: | |
ymin, xmin, ymax, xmax = bbox | |
p1 = (int(xmin * cols), int(ymin * rows)) | |
p2 = (int(xmax * cols), int(ymax * rows)) | |
result[label].append([p1, p2, score]) | |
return result | |
def run_inference_for_single_image(path): | |
image = Image.open(path) | |
# the array based representation of the image will be used later in order to prepare the | |
# result image with boxes and labels on it. | |
image_np = load_image_into_numpy_array(image) | |
# Expand dimensions since the model expects images to have shape: [1, None, None, 3] | |
image_np_expanded = np.expand_dims(image_np, axis=0) | |
image = image_np_expanded | |
with DETECTION_GRAPH.as_default(): | |
with tf.Session() as sess: | |
# Get handles to input and output tensors | |
ops = tf.get_default_graph().get_operations() | |
all_tensor_names = {output.name for op in ops for output in op.outputs} | |
tensor_dict = {} | |
for key in [ | |
'num_detections', 'detection_boxes', 'detection_scores', | |
'detection_classes', 'detection_masks' | |
]: | |
tensor_name = key + ':0' | |
if tensor_name in all_tensor_names: | |
tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(tensor_name) | |
if 'detection_masks' in tensor_dict: | |
# The following processing is only for single image | |
detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0]) | |
detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0]) | |
# Reframe is required to translate mask from box coordinates to image coordinates and fit the image size. | |
real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32) | |
detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1]) | |
detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1]) | |
detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks( | |
detection_masks, detection_boxes, image.shape[1], image.shape[2]) | |
detection_masks_reframed = tf.cast( | |
tf.greater(detection_masks_reframed, 0.5), tf.uint8) | |
# Follow the convention by adding back the batch dimension | |
tensor_dict['detection_masks'] = tf.expand_dims( | |
detection_masks_reframed, 0) | |
image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0') | |
# Run inference | |
output_dict = sess.run(tensor_dict, feed_dict={image_tensor: image}) | |
# all outputs are float32 numpy arrays, so convert types as appropriate | |
output_dict['num_detections'] = int(output_dict['num_detections'][0]) | |
output_dict['detection_classes'] = output_dict['detection_classes'][0].astype(np.int64) | |
output_dict['detection_boxes'] = output_dict['detection_boxes'][0] | |
output_dict['detection_scores'] = output_dict['detection_scores'][0] | |
if 'detection_masks' in output_dict: | |
output_dict['detection_masks'] = output_dict['detection_masks'][0] | |
return output_dict | |
if __name__ == "__main__": | |
test_img_dir = "test/img" | |
test_imgs = os.listdir(test_img_dir) | |
# for path in TEST_IMAGE_PATHS: | |
for fn in test_imgs: | |
path = f"{test_img_dir}/{fn}" | |
img = cv2.imread(path) | |
result = detect(path) | |
for panel in result["solar_panel"]: | |
p1, p2, score = panel | |
cv2.rectangle(img, p1, p2, (0, 255, 0), 2) | |
for hotspot in result["hotspot"]: | |
p1, p2, score = hotspot | |
cv2.rectangle(img, p1, p2, (0, 0, 255), 2) | |
cv2.imshow("result", img) | |
k = cv2.waitKey(0) | |
if k == 13: # enter | |
continue | |
if k == 27: # esc | |
cv2.destroyAllWindows() | |
exit() | |
cv2.destroyAllWindows() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment