行为识别逻辑

拼搏现实的明天。 2022-09-01 08:08 270阅读 0赞

import numpy as np
    import cv2
    import time
    import logging
    
    from actiondetection.api.core.detector import Detector
    from rapid import api as pd
    from deepsort import api as deepsort
    
    
    def get_color_by_id(id):
        colors = [
            (255,0,0),
            (255,255,0),
            (255,0,255),
            (255,255,255),
            (30,40,50),
            (0,255,0),
            (0,50,100),
            (100,0,30),
            (100,100,0),
            (20,10,100),
            (20,250,100),
            (45,0,90),
            (15,10,190),
            (15,100,100),
        ]
        index = id % len(colors)
        return colors[index]
    
    
    class ActionRec(object):
        def __init__(self, people_kwargs, action_kwargs, tracker_kwargs, mall_kwargs,
                     window_frame_count=30, min_scan_count=3, min_pay_count=1,
                     min_stay_duration1=5, min_stay_duration2=20, max_stay_duration2=50):
            """
            :param people_kwargs:行人检测参数
            :param action_kwargs:行为检测识别参数
            :param tracker_kwargs:行人跟踪参数
            :param mall_kwargs:店铺参数
            :param window_frame_count:行为识别窗口长度
            :param min_scan_count:识别为扫码的最小帧数
            :param min_pay_count:识别为支付的最小帧数
            :param min_stay_duration1:关联到扫码/支付行为的id最小停留时间
            :param min_stay_duration2:未关联到扫码/支付行为的id最小停留时间
            :param max_stay_duration2:未关联到扫码/支付行为的id最大停留时间
            """
            self.detect_people_img_size = people_kwargs.pop('img_size', (512, 288))
            logging.info('people detector init')
            kwargs = dict(
                interpolation=cv2.INTER_LINEAR,
            )
            kwargs.update(people_kwargs)
            self.detector_people = pd.Detector(**kwargs)
            logging.info('people detector init end')
    
            logging.info('action detector init')
            kwargs = dict()
            kwargs.update(action_kwargs)
            self.detector_action = Detector(**kwargs)
            logging.info('action detector init end')
    
            logging.info('tracker init')
            kwargs = dict(
                max_iou_distance=0.7,
                max_age=15,
                n_init=3,
                km_f_std_weight_position=1./20,
                km_f_std_weight_velocity=1./160,
            )
            kwargs.update(tracker_kwargs)
            self.tracker = deepsort.IOUTracker(**kwargs)
            logging.info('tracker init end')
    
            self.personarea = np.float32(mall_kwargs['personarea'])
            self.actionarea = np.float32(mall_kwargs['actionarea'])
            self.window_frame_count = window_frame_count
            self.min_scan_count = min_scan_count
            self.min_pay_count = min_pay_count
            self.min_stay_duration1 = min_stay_duration1
            self.min_stay_duration2 = min_stay_duration2
            self.max_stay_duration2 = max_stay_duration2
    
            self.frame_id = 0
            self.alive_tracks_dict = {}
    
        def analyze(self, img, timestamp=None, debug=False):
            if timestamp is None:
                timestamp = time.time()
            self.frame_id += 1
    
            t0 = time.time()
    
            # 行人检测
            dets_people = self.detector_people.detect([img], img_size=self.detect_people_img_size, rotated_nms=False)[0]
    
            # 仅保留在行人区域内的检测结果
            dets_people = [
                det for det in dets_people
                if cv2.pointPolygonTest(self.personarea, ((det['bbox'][0]+det['bbox'][2])/2, det['bbox'][3]), False) >= 0
            ]
    
            t1 = time.time()
    
            # 行人跟踪
            tracker_input = [{
                'tlwh': [det['bbox'][0], det['bbox'][1], det['bbox'][2]-det['bbox'][0], det['bbox'][3]-det['bbox'][1]],
                'score': det['score'],
            } for det in dets_people]
            tracks = self.tracker.track(None, tracker_input)
    
            for track in tracks:
                if track['id'] not in self.alive_tracks_dict:
                    self.alive_tracks_dict[track['id']] = {
                        't0': timestamp,
                        'scan': False,
                        'pay': False,
                        'scan_cache': [],
                        'pay_cache': [],
                    }
    
            t2 = time.time()
    
            # 动作检测
            dets = self.detector_action.detect([img])[0]
    
            # 仅保留在动作区域内的检测结果
            dets = [
                det for det in dets
                if cv2.pointPolygonTest(self.actionarea, ((det[0]+det[2])/2, det[3]), False) >= 0
            ]
    
            if len(tracks) > 0:
                centers_people = np.array([((x['bbox'][0]+x['bbox'][2])/2, (x['bbox'][1]+x['bbox'][3])/2) for x in tracks])
                for det in dets:
                    # 将扫商品/支付动作关联到离它距离最近的顾客
                    (x1, y1, x2, y2), cls = det[0:4], round(det[-1])
                    center = ((x1+x2)/2, (y1+y2)/2)
                    distances = np.linalg.norm(centers_people-center, axis=1)
                    i = np.argmin(distances)
                    td = self.alive_tracks_dict[tracks[i]['id']]
    
                    # 滑动窗口的方式判断是否满足扫商品/支付的条件
                    cache = td['scan_cache'] if cls == 0 else td['pay_cache']
                    if len(cache) == 0 or cache[-1] != self.frame_id:
                        cache.append(self.frame_id)
                    while len(cache) > 0 and self.frame_id-cache[0] > self.window_frame_count-1:
                        del cache[0]
                    if cls == 0 and len(cache) >= self.min_scan_count:
                        td['scan'] = True
                    elif cls == 1 and len(cache) >= self.min_pay_count:
                        td['pay'] = True
    
            # 根据消失的id产生交易记录
            reserved = {track['id']: True for track in tracks}
            results = []
            for k, v in self.alive_tracks_dict.items():
                if k not in reserved:
                    duration = timestamp-v['t0']
                    if (v['scan'] or v['pay']) and duration >= self.min_stay_duration1 \
                            or self.min_stay_duration2 <= duration <= self.max_stay_duration2:
                        results.append({
                            'timestamp': v['t0'],
                            'scan': v['scan'],
                            'pay': v['pay'],
                        })
            self.alive_tracks_dict = {k: v for k, v in self.alive_tracks_dict.items() if k in reserved}
    
            if debug:
                for track in tracks:
                    x1, y1, x2, y2 = map(int, track['bbox'])
                    color = get_color_by_id(track['id'])
                    cv2.rectangle(img, (x1, y1), (x2, y2), color, 4, lineType=cv2.LINE_AA)
                    text = '%d' % track['id']
                    if self.alive_tracks_dict[track['id']]['scan']:
                        text += ', scan'
                    if self.alive_tracks_dict[track['id']]['pay']:
                        text += ', pay'
                    cv2.putText(img, text, (x1+10, y2-15), cv2.FONT_HERSHEY_SIMPLEX, 1.2, color, 4)
                for det in dets:
                    (x1, y1, x2, y2), cls = map(int, det[0:4]), round(det[-1])
                    color = [(0, 255, 0), (255, 255, 0)][cls]
                    cv2.rectangle(img, (x1, y1), (x2, y2), color, 4, lineType=cv2.LINE_AA)
                    text = ['scan', 'pay'][cls]
                    cv2.putText(img, text, (x1+10, y2-15), cv2.FONT_HERSHEY_SIMPLEX, 1.2, color, 4)
    
            t3 = time.time()
    
            logging.debug('[scan_pay_action] people_detect: %.4fs | people_track: %.4fs | action_rec: %.4fs | total: %.4fs' %
                          (t1-t0, t2-t1, t3-t2, t3-t0))
    
            return results