src/visualize.py

import numpy as np
from PIL import Image, ImageDraw, ImageFile
from .NetWork import VGG
import paddle
import cv2

def get_color_map_list(num_classes):
    """
    Args:
        num_classes (int): number of class
    Returns:
        color_map (list): RGB color list
    """
    color_map = num_classes * [0, 0, 0]
    for i in range(0, num_classes):
        j = 0
        lab = i
        while lab:
            color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
            color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
            color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
            j += 1
            lab >>= 3
    color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
    return color_map


def draw_det(image, dt_bboxes, name_set):
    im = Image.fromarray(image)
    draw_thickness = min(im.size) // 320
    draw = ImageDraw.Draw(im)
    clsid2color = {}
    color_list = get_color_map_list(len(name_set))

    for (cls_id, score, xmin, ymin, xmax, ymax) in dt_bboxes:
        image_box = im.crop(tuple([xmin, ymin, xmax, ymax]))
        label = emotic(image_box)
        cls_id = int(cls_id)
        color = tuple(color_list[cls_id])
        # draw bbox
        draw.line(
            [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
             (xmin, ymin)],
            width=draw_thickness,
            fill=color)

        # draw label
        text = "{} {:.4f}".format(label, score)
        box = draw.textbbox((xmin, ymin), text, anchor='lt')
        draw.rectangle(box, fill=color)
        draw.text((box[0], box[1]), text, fill=(255, 255, 255))
    image = np.array(im)
    return image


def emotic(image):
    def load_image(img):
        # 将图片尺寸缩放道 224x224
        img = cv2.resize(img, (224, 224))
        # 读入的图像数据格式是[H, W, C]
        # 使用转置操作将其变成[C, H, W]
        img = np.transpose(img, (2, 0, 1))
        img = img.astype('float32')
        # 将数据范围调整到[-1.0, 1.0]之间
        img = img / 255.
        img = img * 2.0 - 1.0
        return img

    model = VGG(num_class=7)
    params_file_path = r'configs/vgg.pdparams'
    img = np.array(image)
    # plt.imshow(img)
    # plt.axis('off')
    # plt.show()

    param_dict = paddle.load(params_file_path)
    model.load_dict(param_dict)
    # 灌入数据
    # model.eval()
    tensor_img = load_image(img)
    tensor_img = np.expand_dims(tensor_img, 0)

    results = model(paddle.to_tensor(tensor_img))
    # 取概率最大的标签作为预测输出
    lab = np.argsort(results.numpy())
    tap = lab[0][-1]

    if tap == 0:
        return 'SAD'
    elif tap == 1:
        return 'DISGUST'
    elif tap == 2:
        return 'HAPPY'
    elif tap == 3:
        return 'FEAR'
    elif tap == 4:
        return 'SUPERISE'
    elif tap == 5:
        return 'NATUREAL'
    elif tap == 6:
        return 'ANGRY'
    else:
        raise ('Not excepted file name')