整合后的代码如下
import sys
import argparse
from yolo import YOLO, detect_video
from PIL import Image
import os
import numpy as np
import six.moves.urllib as urllib
import tarfile
import tensorflow as tf
import zipfile
from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
sys.path.append("..")
from utils import label_map_util
from utils import visualization_utils as vis_util
import cv2
from timeit import default_timer as timer
PATH_TO_CKPT = 'D:\\Faster_RCNN_\\output'+'\\frozen_inference_graph.pb'
PATH_TO_LABELS = 'D:/Faster_RCNN_/dataset/pascal_label_map.pbtxt'
NUM_CLASSES = 3
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)
print(category_index)
'''
def load_image_into_numpy_array(image):
#image = Image.merge('RGB',(image,image,image))
img = np.array(image)
#print(image.size)
return img
'''
IMAGE_SIZE = (12, 8)
with detection_graph.as_default():
with tf.Session(graph=detection_graph) as sess:
while 1:
# for image_path in os.listdir('test_images_voc2'):
# for image_path in os.listdir('./test_pointer'):
image = Image.open('D:\\models-master\\models-master\\research\\object_detection2\\test_pointer1\\test.jpg' )
r_image = YOLO().detect_image(image)
pilimg = Image.fromarray(r_image)
pilimg.save('D:\\models-master\\models-master\\research\\object_detection2\\result_pointer1\\result.jpg')
print('success')
start = timer()
s_image = Image.open('D:\\models-master\\models-master\\research\\object_detection2\\result_pointer1\\result.jpg')
image_np = np.array(s_image)
#image_np = np.array(r_image)
#image_np_expanded = np.expand_dims(image_np, axis=0)
image_np_expanded = np.expand_dims(image_np, axis=0)
#image233 = np.concatenate(image_np_expanded, axis=-1)
print(s_image.size)
image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
scores = detection_graph.get_tensor_by_name('detection_scores:0')
classes = detection_graph.get_tensor_by_name('detection_classes:0')
num_detections = detection_graph.get_tensor_by_name('num_detections:0')
(boxes, scores, classes, num_detections) = sess.run(
[boxes, scores, classes, num_detections],
feed_dict={image_tensor: image_np_expanded})
end = timer()
vis_util.visualize_boxes_and_labels_on_image_array(
image_np, np.squeeze(boxes),
np.squeeze(classes).astype(np.int32),
np.squeeze(scores),
category_index,
use_normalized_coordinates=True,
line_thickness=8)
(r, g, b)=cv2.split(image_np)
image_np=cv2.merge([b,g,r])
print(end - start)
# cv2.imwrite('./result_images_voc2'+'/'+image_path,image_np)
cv2.imwrite('D:\\models-master\\models-master\\research\\object_detection2\\result_pointer1\\result1.jpg', image_np)
重点是这句话pilimg = Image.fromarray(r_image)
,必须要这句话不然会报错AttributeError: ‘numpy.ndarray’ object has no attribute 'save’原因是下面的yolo代码的最后 image1 = cv2.cvtColor(np.array(image), cv2.COLOR_GRAY2BGR)
是opencv下的图片为numpy格式,而上述代码是pillow下的,所以要做一下转换。
同理 image1 = cv2.cvtColor(np.array(image), cv2.COLOR_GRAY2BGR)中的第一个参数若是image也会报错,要从pillow格式下转换为opencv的numpy格式。
yolo里的detect_image(self, image)代码如下
def detect_image(self, image):
start = timer()
if self.model_image_size != (None, None):
assert self.model_image_size[0]%32 == 0, 'Multiples of 32 required'
assert self.model_image_size[1]%32 == 0, 'Multiples of 32 required'
boxed_image = letterbox_image(image, tuple(reversed(self.model_image_size)))
else:
new_image_size = (image.width - (image.width % 32),
image.height - (image.height % 32))
boxed_image = letterbox_image(image, new_image_size)
image_data = np.array(boxed_image, dtype='float32')
#print(image_data.shape)
image_data /= 255.
image_data = np.expand_dims(image_data, 0) # Add batch dimension.
out_boxes, out_scores, out_classes = self.sess.run(
[self.boxes, self.scores, self.classes],
feed_dict={
self.yolo_model.input: image_data,
self.input_image_shape: [image.size[1], image.size[0]],
K.learning_phase(): 0
})
print('Found {} boxes for {}'.format(len(out_boxes), 'img'))
font = ImageFont.truetype(font='font/FiraMono-Medium.otf',
size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
thickness = (image.size[0] + image.size[1]) // 300
for i, c in reversed(list(enumerate(out_classes))):
predicted_class = self.class_names[c]
box = out_boxes[i]
score = out_scores[i]
label = '{} {:.2f}'.format(predicted_class, score)
draw = ImageDraw.Draw(image)
#print(image.size)
label_size = draw.textsize(label, font)
top, left, bottom, right = box
top = max(0, np.floor(top + 0.5).astype('int32'))
left = max(0, np.floor(left + 0.5).astype('int32'))
bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
print(label, (left, top), (right, bottom))
if top - label_size[1] >= 0:
text_origin = np.array([left, top - label_size[1]])
else:
text_origin = np.array([left, top + 1])
# My kingdom for a good redistributable image drawing library.
for i in range(thickness):
draw.rectangle(
[left + i, top + i, right - i, bottom - i],
outline=255) #outline=self.colors[c]
draw.rectangle(
[tuple(text_origin), tuple(text_origin + label_size)],
fill=255) #fill=self.colors[c]
draw.text(text_origin, label, fill=0, font=font)# fill=255
del draw
#img = cv2.imdecode(image, cv2.IMREAD_GRAYSCALE)
#img = cv2.imread(image)
#cropped = img[bottom:top, left:right]
image = image.crop((left, top, right, bottom))
image1 = cv2.cvtColor(np.array(image), cv2.COLOR_GRAY2BGR)
#print(len(image.split(image)))
end = timer()
print(end - start)
return image1
#return cropped