#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue May 29 10:53:04 2018
@author: sisyphus
"""
import os
import random
import json
import colorsys
import numpy as np
import tensorflow as tf
from keras import backend as K
from keras.models import load_model
from PIL import Image, ImageFont, ImageDraw
from timeit import time
def sigmoid(x):
y = 1/(1+np.exp(-x))
return y
def letterbox_image(image, size):
'''resize image with unchanged aspect ratio using padding'''
image_w, image_h = image.size
w, h = size
new_w = int(image_w * min(w*1.0/image_w, h*1.0/image_h))
new_h = int(image_h * min(w*1.0/image_w, h*1.0/image_h))
resized_image = image.resize((new_w,new_h), Image.BICUBIC)
boxed_image = Image.new('RGB', size, (128,128,128))
boxed_image.paste(resized_image, ((w-new_w)//2,(h-new_h)//2))
return boxed_image
def yolo_head(feats, anchors, num_classes, input_shape):
"""Convert final layer features to bounding box parameters."""
num_anchors = len(anchors)
anchors_tensor = np.reshape(anchors, [1, 1, 1, num_anchors, 2])#++
grid_shape = np.shape(feats)[1:3]#++
grid_y = np.tile(np.reshape(np.arange(0, stop=grid_shape[0]), [-1, 1, 1, 1]),
[1, grid_shape[1], 1, 1])#++
grid_x = np.tile(np.reshape(np.arange(0, stop=grid_shape[1]), [1, -1, 1, 1]),
[grid_shape[0], 1, 1, 1])#++
grid = np.concatenate([grid_x, grid_y], axis=-1)#++
grid = grid.astype(np.float32)#++
feats = np.reshape(
feats, [-1, grid_shape[0], grid_shape[1], num_anchors, num_classes + 5])
box_xy = sigmoid(feats[..., :2])
box_wh = np.exp(feats[..., 2:4])
box_confidence = sigmoid(feats[..., 4:5])
box_class_probs = sigmoid(feats[..., 5:])
box_xy = (box_xy + grid) / grid_shape[::-1]
box_wh = box_wh * anchors_tensor / input_shape[::-1]
return box_xy, box_wh, box_confidence, box_class_probs
def yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape):
'''Get corrected boxes'''
box_yx = box_xy[..., ::-1]
box_hw = box_wh[..., ::-1]
mm = min(input_shape[0]/image_shape[0],input_shape[1]/image_shape[1])#++
image_shape = [image_shape[0]*mm, image_shape[1]*mm]#++
new_shape = [round(image_shape[0]),round(image_shape[1])]#++
offset = (input_shape[0]-new_shape[0])/2./input_shape[0]
scale = input_shape[0]/new_shape[0]
box_yx = (box_yx - offset) * scale
box_hw *= scale
box_mins = box_yx - (box_hw / 2.)
box_maxes = box_yx + (box_hw / 2.)
boxes = np.concatenate([
box_mins[..., 0:1], # y_min
box_mins[..., 1:2], # x_min
box_maxes[..., 0:1], # y_max
box_maxes[..., 1:2] # x_max
],axis = -1)
# Scale boxes back to original image shape.
boxes *= np.concatenate([image_shape, image_shape],axis = -1)
return boxes
def yolo_boxes_and_scores(feats, anchors, num_classes, input_shape, image_shape):
'''Process Conv layer output'''
box_xy, box_wh, box_confidence, box_class_probs = yolo_head(feats,
anchors, num_classes, input_shape)
boxes = yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape)
boxes = np.reshape(boxes, [-1, 4])
box_scores = box_confidence * box_class_probs
box_scores = np.reshape(box_scores, [-1, num_classes])
return boxes, box_scores
def nms(bounding_boxes, confidence_score, threshold):
# If no bounding boxes, return empty list
if len(bounding_boxes) == 0:
return [], []
# coordinates of bounding boxes
start_x = bounding_boxes[:, 0]
start_y = bounding_boxes[:, 1]
end_x = bounding_boxes[:, 2]
end_y = bounding_boxes[:, 3]
# Picked bounding boxes
picked_boxes = []
picked_score = []
index_b = []
# Compute areas of bounding boxes
areas = (end_x - start_x + 1) * (end_y - start_y + 1)
# Sort by confidence score of bounding boxes
score = np.array(confidence_score)
order = np.argsort(score)
# Iterate bounding boxes
while order.size > 0:
# The index of largest confidence score
index = order[-1]
# Pick the bounding box with largest confidence score
picked_boxes.append(bounding_boxes[index])
picked_score.append(confidence_score[index])
index_b.append(index)
# Compute ordinates of intersection-over-union(IOU)
x1 = np.maximum(start_x[index], start_x[order[:-1]])
x2 = np.minimum(end_x[index], end_x[order[:-1]])
y1 = np.maximum(start_y[index], start_y[order[:-1]])
y2 = np.minimum(end_y[index], end_y[order[:-1]])
# Compute areas of intersection-over-union
w = np.maximum(0.0, x2 - x1 + 1)
h = np.maximum(0.0, y2 - y1 + 1)
intersection = w * h
# Compute the ratio between intersection and union
ratio = intersection / (areas[index] + areas[order[:-1]] - intersection)
left = np.where(ratio < threshold)
order = order[left]
return index_b
def boolean_mask(N,M):
ax = N.shape[0]
id=[]
for i in range(0,ax):
j=i*M[i]
id.append(j)
boxes=[]
for k in range(0,ax):
if id[k]!=0:
box=N[k]
boxes.append(box)
return np.array(boxes)
def yolo_eval(yolo_outputs,
anchors,
num_classes,
image_shape,
max_boxes=20,
score_threshold=.6,
iou_threshold=.5):
"""Evaluate YOLO model on given input and return filtered boxes."""
anchor_mask = [[6,7,8], [3,4,5], [0,1,2]]
input_shape = np.shape(yolo_outputs[0])[1]*32, np.shape(yolo_outputs[0])[2]*32 #++
boxes = []
box_scores = []
for l in range(3):
_boxes, _box_scores = yolo_boxes_and_scores(yolo_outputs[l],
anchors[anchor_mask[l]], num_classes, input_shape, image_shape)
boxes.append(_boxes)
box_scores.append(_box_scores)
boxes = np.concatenate(boxes, axis=0)
boxes = boxes.astype(np.float32)#++
box_scores = np.concatenate(box_scores, axis=0)
mask = box_scores >= score_threshold
boxes_ = []
scores_ = []
classes_ = []
for c in range(num_classes):
class_boxes = boolean_mask(boxes, mask[:, c])
if np.shape(class_boxes)[0] == 0:
continue
class_box_scores = boolean_mask(box_scores[:, c], mask[:, c])
nms_index = nms(class_boxes, class_box_scores, iou_threshold)#++
class_boxestt_=[]
class_box_scorestt_=[]
classestt_=[]
for i in nms_index:
class_boxestt = class_boxes[i]
class_box_scorestt = class_box_scores[i]
classestt = c
class_boxestt_.append(class_boxestt)
class_box_scorestt_.append(class_box_scorestt)
classestt_.append(classestt)
boxes_.append(class_boxestt_)
scores_.append(class_box_scorestt_)
classes_.append(classestt_)
if np.shape(boxes_)[0] == 0:
return [],[],[]
else:
# return boxes_[0], scores_[0], classes_[0]
return boxes_, scores_, classes_
def get_anchors(anchors_path):
'''
获取anchors
'''
anchors_path = os.path.expanduser(anchors_path)
with open(anchors_path) as f:
anchors = f.readline()
anchors = [float(x) for x in anchors.split(',')]
anchors = np.array(anchors).reshape(-1, 2)
return anchors
def get_class(classes_path):
'''
获取检测类别
'''
classes_path = os.path.expanduser(classes_path)
with open(classes_path) as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
def postprocess(image, result):
'''
后处理函数
输入：待检测图片image，yolo网络输出result（3组向量，对应3组anchors）
输出：检测信息json(保存在文件夹)，带标注框带图片image(作为函数返回值)
'''
start = time.time()
anchors_path = 'model_data/yolo_anchors.txt'
classes_path = 'model_data/voc_6c_classes.txt'
score = 0.3
iou = 0.5
input_image_shape = [image.size[1], image.size[0]]
# anchors = get_anchors(anchors_path)
anchors=np.array([[ 10., 13.],
[ 16., 30.],
[ 33., 23.],
[ 30., 61.],
[ 62., 45.],
[ 59., 119.],
[116., 90.],
[156., 198.],
[373., 326.]])
# class_names = get_class(classes_path)
class_names = ['crazing', 'inclusion', 'patches', 'pitted_surface', 'rolled-in_scale', 'scratches']
out_boxes, out_scores, out_classes = yolo_eval(result,
anchors,
len(class_names),
input_image_shape,
score_threshold=score,
iou_threshold=iou)
out_boxes, out_scores, out_classes = sum(out_boxes, []),sum(out_scores, []),sum(out_classes, [])
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
# Generate colors for drawing bounding boxes.
hsv_tuples = [(x / len(class_names), 1., 1.)
for x in range(len(class_names))]
colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
colors = list(
map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
colors))
random.seed(10101) # Fixed seed for consistent colors across runs.
random.shuffle(colors) # Shuffle colors to decorrelate adjacent classes.
random.seed(None) # Reset seed to default.
resultsForJSON = []#+++
for i, c in reversed(list(enumerate(out_classes))):
predicted_class = class_names[c]
box = out_boxes[i]
score = out_scores[i]
label = '{} {:.2f}'.format(predicted_class, score)
draw = ImageDraw.Draw(image)
label_size = draw.textsize(label, font)
top, left, bottom, right = box/2 #++
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))
area=(bottom-top)*(right-left)#####+
resultsForJSON.append({"label": predicted_class,\
"confidence": float('%.2f' % score),\
"topleft": {"x": float(left), "y": float(top)},\
"bottomright": {"x": float(right), "y": float(bottom)},\
"area":float(area)})
if top - label_size[1] >= 0:
text_origin = np.array([left, top - label_size[1]])
else:
text_origin = np.array([left, top + 1])
for i in range(thickness):
draw.rectangle(
[left + i, top + i, right - i, bottom - i],
outline=colors[c])
draw.rectangle(
[tuple(text_origin), tuple(text_origin + label_size)],
fill=colors[c])
draw.text(text_origin, label, fill=(0, 0, 0), font=font)
del draw
end = time.time()
print(end - start)
####################保存json##########start
print(resultsForJSON)
imgdir = "/Users/sisyphus/keras_yolo3_model/" ###########+路径也要修改
outfolder = os.path.join(imgdir, 'out')
img_name = os.path.join(outfolder, os.path.basename(pathall))
if resultsForJSON == []:
resultsForJSON.append({"label": "Normal",\
"confidence": float(0), \
"topleft": {"x": float(0), "y": float(0)}, \
"bottomright": {"x": float(0), "y": float(0)},\
"area":float(0)})
textJSON = json.dumps(resultsForJSON)
textFile = os.path.splitext(img_name)[0] + ".json"
print('\n')
with open(textFile, 'w') as f:
f.write(textJSON)
print('Normal\n')
else:
textJSON = json.dumps(resultsForJSON)
textFile = os.path.splitext(img_name)[0] + ".json"
print(textFile)
print('\n')
with open(textFile, 'w') as f:
f.write(textJSON)
####################保存json##########end
return image
imgdir = "/Users/sisyphus/keras_yolo3_model/" #模型存放路径
outfolder = os.path.join(imgdir, 'out') #处理后图片和json保存路径
path = "/Users/sisyphus/keras_yolo3_model/sample_img_test2" #待处理图片存放路径
model_path = 'model_data/yolo.h5'
print("Loading model......")
yolo_model = load_model(model_path, compile=False) #加载模型
print("Model loaded successful")
for xd in os.listdir(path):
# xd = os.listdir(path)[0]
pathall = os.path.join(path,xd)
image = Image.open(pathall)
boxed_image = letterbox_image(image, tuple(reversed((416, 416))))
image_data = np.array(boxed_image, dtype='float32')
print(image_data.shape)
image_data /= 255.
image_data = np.expand_dims(image_data, 0)
# input_image_shape = K.placeholder(shape=(2, ))
result = yolo_model.predict(image_data) #模型推导
# result = result[0]
imgpost = postprocess(image, result) #后处理函数
img_name = os.path.join(outfolder, os.path.basename(pathall))
imgpost.save(img_name) #保存处理后的图片
