Spaces:
Paused
Paused
File size: 2,954 Bytes
12645de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
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')
|