【2022ModelBox客流分析活动打卡】作业1

255 2023-07-17

track 255 bbox img_data 170

题目描述

1、修改课程1的显示效果，不同人的检测框显示不同的颜色，同时显示出运动轨迹（200分）

完成要求：需要贴出改进后的核心代码和程序运行效果截图，效果如下：

核心代码

draw_passenger_bbox.py 文件

workspace\passenger_flow_count\etc\flowunit\draw_passenger_bbox\draw_passenger_bbox.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

          
# Copyright (C) 2022 Huawei Technologies Co., Ltd. All rights reserved.

          
import _flowunit as modelbox
import cv2
import json
import numpy as np
from PIL import Image, ImageDraw, ImageFont

          

          
class draw_passenger_bboxFlowUnit(modelbox.FlowUnit):
def __init__(self):
super().__init__()
        
          self.colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0],
[0, 255, 0], [0, 255, 85], [0, 255, 170],
[0, 0, 255], [85, 0, 255], [170, 0, 255],
[0, 85, 255], [255, 0, 170], [255, 0, 85],
[170, 255, 0], [85, 255, 0], [0, 170, 255],
[0, 255, 255], [255, 255, 0], [255, 0, 255]]
        
          self.num_colors = len(self.colors)
        
          self.points_list = []

          

          
def open(self, config):
# 获取功能单元的配置参数
        
          self.area = config.get_int_list('area', [])
        
          self.area = np.array(self.area).reshape(-1, 1, 2).astype(np.int32)
        
          self.font_path = config.get_string("font_path")
return modelbox.Status.StatusCode.STATUS_SUCCESS

          
def process(self, data_context):
# 从DataContext中获取输入输出BufferList对象
        
          in_image = data_context.input("in_image")
        
          in_track = data_context.input("in_track")
        
          out_image = data_context.output("out_image")

          
# 循环处理每一个输入Buffer数据
for buffer_img, buffer_track in zip(in_image, in_track):
# 输入图像Buffer获取宽、高、通道数等属性信息
            
          width = buffer_img.get('width')
            
          height = buffer_img.get('height')
            
          channel = buffer_img.get('channel')

          
# 将输入Buffer转换为Python对象
            
          img_data = np.array(buffer_img.as_object(), copy=False)
            
          img_data = img_data.reshape((height, width, channel))

          
# 获取json字符串格式的跟踪结果
            
          track_json = buffer_track.as_object()

          
# 业务处理：从前一个功能单元的json数据中获取检测框，画在图像上
            
          tracking_objects = self.decode_track_info(track_json, (height, width))
            
          img_out = self.draw_tracking_object(img_data, tracking_objects)

          
# 将画框后的图像转换为Buffer
            
          out_buffer = modelbox.Buffer(self.get_bind_device(), img_out)

          
# 设置输出Buffer的Meta信息，此处直接拷贝输入Buffer的Meta信息
            
          out_buffer.copy_meta(buffer_img)

          
# 将输出Buffer放入输出BufferList中
            
          out_image.push_back(out_buffer)

          
# 返回成功标志，ModelBox框架会将数据发送到后续的功能单元
return modelbox.Status.StatusCode.STATUS_SUCCESS

          
def decode_track_info(self, track_json, input_shape):
'''从json数据中解码出检测框'''
try:
            
          tracking_objects = []
            
          track_result = json.loads(track_json)
            
          tracking_list = json.loads(track_result['tracking_objects'])
for track_dict in tracking_list:
                
          tracking_obj = {}
                
          tracking_obj["id"] = track_dict["id"]
                
          x1 = int(track_dict["bbox"][0] * input_shape[1])
                
          y1 = int(track_dict["bbox"][1] * input_shape[0])
                
          x2 = int(track_dict["bbox"][2] * input_shape[1])
                
          y2 = int(track_dict["bbox"][3] * input_shape[0])
                
          tracking_obj["bbox"] = [x1, y1, x2, y2]
                
          tracking_objects.append(tracking_obj)
except Exception as ex:
            
          modelbox.error(str(ex))
return None, None
else:
return tracking_objects

          
def draw_tracking_object(self, img_data, tracking_objects):
'''在图中画出跟踪对象的检测框和过线的行人数据'''

          
        
          thickness = 2
        
          GRAY = (117, 117, 117)
        
          GREEN = (0, 255, 0)
        
          YELLO = (255, 255, 0)
        
          points = []

          
# 画出区域边界线
        
          cv2.polylines(img_data, [self.area], True, YELLO, 3)

          
        
          flow_count = 0
for track in tracking_objects:
# 人形框的中心点
            
          c_x = int((track["bbox"][0] + track["bbox"][2]) / 2)
            
          c_y = int((track["bbox"][1] + track["bbox"][3]) / 2)

          

          
#中心点画点
            
          cv2.line(img_data,(c_x, c_y),(c_x,c_y),GREEN,5)
            
          track_id = track["id"]
            
          color = self.colors[track_id % self.num_colors]
# print(track.type)

          

          
# 判断人形框的中心点是否在区域内
            
          flag = cv2.pointPolygonTest(self.area, (c_x, c_y), False)
if flag > 0:
# 区域内人形框用绿色，同时客流计数增加
                
          flow_count += 1
                
          cv2.rectangle(img_data, (track["bbox"][0], track["bbox"][1]),
(track["bbox"][2], track["bbox"][3]), color, 2)

          
                
          points.append(((c_x, c_y),color))
                
          self.points_list.append(points)

          
#只保存最近60个点
if(len(self.points_list) > 60):
                    
          self.points_list.pop(0)

          
#画出跟踪点
for pp in self.points_list:
for point,cc in pp:
                         
          cv2.circle(img_data,point,3,cc,6)
#cv2.line(img_data,point,point,cc,2)
else:
# 区域内人形框用灰色
                
          cv2.rectangle(img_data, (track["bbox"][0], track["bbox"][1]),
(track["bbox"][2], track["bbox"][3]), GRAY, thickness)

          
# 左上角显示实时的客流数量
        
          img_data = self.put_chi_text(
            
          img_data, 'JaneConanの作业 客流计数：%d 人' % flow_count, (50, 20), YELLO, 50)

          
return img_data

          
def put_chi_text(self, img, text, location, color=(0, 255, 0), size=50):
'''在图片中写汉字'''
if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
            
          img = Image.fromarray(img)

          
        
          draw = ImageDraw.Draw(img)  # 创建一个可以在给定图像上绘图的对象
        
          font_style = ImageFont.truetype(
            
          self.font_path, size, encoding="utf-8")  # 字体的格式
        
          draw.text(location, text, color, font=font_style)  # 绘制文本
return np.asarray(img)  # 转换回OpenCV格式

          
def close(self):
return modelbox.Status()