keras_yolov3_test_delk

#!/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)       #保存处理后的图片