多行人計(jì)數(shù)與跟蹤

from utils.datasets import *from utils.utils import *import osimport torchfrom models.create_model import Create_Model
os.environ['CUDA_VISIBLE_DEVICES'] = "0"print(torch.cuda.is_available())

# 獲得特征距離def person_distance(person_encodings, person_unknow):    if len(person_encodings) == 0:        return np.empty((0))    l1 = np.sqrt(np.sum(np.square(person_encodings - person_unknow), axis=-1))    return l1
#判斷閾值def com_person(person_list, person, tolerance=1):    dis = person_distance(person_list, person)    # print(dis)    return dis,list(dis <= tolerance)

# 組合距離矩陣def get_iou(boxes1,boxes2,arr_frame):    boxes1 = np.array(boxes1,dtype=np.float32)    boxes2 = np.array(boxes2,dtype=np.float32)
    # 求左上角右下角坐標(biāo)    b1 = np.expand_dims(boxes1, -2)    b1_xy = b1[..., :2]    b1_wh = b1[..., 2:4]    b1_half = b1_wh / 2.    b1_mins = b1_xy - b1_half    b1_maxs = b1_xy + b1_half
    b2 = np.expand_dims(boxes2, 0)    b2_xy = b2[..., :2]    b2_wh = b2[..., 2:4]    b2_half = b2_wh / 2.    b2_mins = b2_xy - b2_half    b2_maxs = b2_xy + b2_half
    # 求交集面積    intersction_min = np.maximum(b1_mins, b2_mins)    intersction_max = np.minimum(b1_maxs, b2_maxs)    intersction_wh = np.maximum(intersction_max - intersction_min, 0.)    intersction_area = intersction_wh[..., 0] * intersction_wh[..., 1]
    # 求交并比    b1_area = b1_wh[..., 0] * b1_wh[..., 1]    b2_area = b2_wh[..., 0] * b2_wh[..., 1]
    # 理論上  交并比越大說明 越接近，不過為了 和 特征距離統(tǒng)一，這里使用1-，確保距離矩陣只拿最小值即可    iou = 1-intersction_area / (b1_area + b2_area - intersction_area)
    # 當(dāng)1-交并比 過大 ，直接設(shè)置一個(gè)很大的數(shù)，方便過濾    iou[iou>0.7]=10086+1e-5
    # 調(diào)整特征矩陣的維度    arr_frame=arr_frame.reshape(arr_frame.shape[0],arr_frame.shape[1])    # 將 iou矩陣與特征矩陣相加 得到最終的距離矩陣    iou = iou+arr_frame    return  iou

def run():    # 加載目標(biāo)檢測模型    device = torch_utils.select_device('cpu')    # google_utils.attempt_download(model_path)    model = torch.load(model_path, map_location=device)['model']    model.to(device).eval()    names = model.names if hasattr(model, 'names') else model.modules.names
    # 加載reid 識(shí)別模型    input_size = (215, 90, 3)    model_, pred_model = Create_Model(inpt=input_size, num_classes=1812)    model_.load_weights('weights\ep039-loss0.066.h5')

    video_capture = cv2.VideoCapture(Cam_num)
    # 寫入視頻    video_FourCC = int(video_capture.get(cv2.CAP_PROP_FOURCC))    video_fps = video_capture.get(cv2.CAP_PROP_FPS)    video_size = (int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)),                  int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))    out = cv2.VideoWriter('output.mp4', video_FourCC, video_fps, video_size)    index = 0
    # 行人特征  行人坐標(biāo)    unknow_person_emb =[]    xyxy_all = []
    # person_state = []
    while True:        ret, im0 = video_capture.read()        # im0 = cv2.flip(im0, 1, dst=None)        iimage = im0.copy()        img = letterbox(im0, new_shape=image_size)[0]        img = img[:, :, ::-1].transpose(2, 0, 1)        img = np.ascontiguousarray(img)        img = torch.from_numpy(img).to(device)        img = img.half() if half else img.float()  # uint8 to fp16/32        img /= 255.0  # 0 - 255 to 0.0 - 1.0        if img.ndimension() == 3:            img = img.unsqueeze(0)        # print(img.shape)
        # Inference        pred = model(img, augment=False)[0]

        if half:            pred = pred.float()
        # Apply NMS        pred = non_max_suppression(pred, conf_thres, iou_thres,                                   fast=True, classes=None, agnostic=False)
        this_frame_xyxy =[]        this_frame_emb=[]
        for i, det in enumerate(pred):            if det is not None and len(det):                # Rescale boxes from img_size to im0 size                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()                for *xyxy, conf, cls in det:                    if names[int(cls)] == 'person':                        c1, c2 = (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3]))                        # 獲得ROI 區(qū)域                        x_min, y_min = c1                        x_max, y_max = c2                        roi = iimage[y_min:y_max, x_min:x_max]                        roi = cv2.resize(roi,(90,215))                        # cv2.imwrite('person.jpg',roi)                        image_1 = np.asarray(roi).astype(np.float64) / 255                        photo1 = np.expand_dims(image_1 ,0)                        output1 = pred_model.predict(photo1)                        # centerx = x_min+(x_max-x_min)/2                        # centery = y_min+(y_max-y_min)/2                        # w , h  =  (x_max-x_min) , (y_max-y_min)
                        # if 0.2<w/h<0.5:                        # 提取每一幀的行人                        this_frame_xyxy.append([x_min,y_min,x_max,y_max])                        this_frame_emb.append(output1)

                # 第一幀只進(jìn)行 特征提取                if index == 0:                    for index_, xyxy_ in enumerate(this_frame_xyxy):
                        unknow_person_emb.append(this_frame_emb[index_])                        xyxy_all.append(xyxy_)
                # 開始特征匹配                else:                    # 當(dāng)前幀的每個(gè)行人 與 上一幀的行人的 特征距離                    this_frame_emb_arr=[]                    for i in this_frame_emb:                        dit, com_p = com_person(unknow_person_emb, i, tolerance=0.7)                        # print(dit)                        this_frame_emb_arr.append(dit)                    arr_frame = np.array(this_frame_emb_arr)                    # 獲得距離矩陣 （物理距離+特征距離）                    ious= get_iou(this_frame_xyxy, xyxy_all,arr_frame)                    # 當(dāng)前幀已有的行人id  確保不出現(xiàn)重復(fù)id                    this_person_index_no_state=[]
                    # 當(dāng)前幀索引，                    for index_,iou in enumerate(ious):                        # 查找距離最近的3個(gè)行人 （物理距離+特征距離）                        ind = np.argpartition(iou, (0,3))[:3]                        # ind = [iou.argmin()]                        # 最小的 距離小于1.4  特征距離0.7 + 物理距離0.7                        if iou[ind[0]]<1.4:                            person_index = ind[0]                            i= 0                            match_person = True                            # 當(dāng) 行人id 已經(jīng)出現(xiàn)時(shí)，取第二 第三個(gè)id 進(jìn)行匹配                            while person_index in this_person_index_no_state:                                person_index =ind[i+1]                                # 匹配失敗則認(rèn)為該行人 是新出現(xiàn)的行人                                if iou[person_index]>1.4:                                    match_person = False                                if i == 1:                                    break                                i+=1
                            # 匹配行人成功   畫框、更新特征以及坐標(biāo)                            if match_person:                            # while iou.argmin() not in this_person_index_no_state:                                x1, y1, x2, y2 = this_frame_xyxy[index_]                                person_id = 'ID-%s' % person_index                                cv2.rectangle(im0, (x1, y1), (x2, y2), (0, 255, 255), 3)                                cv2.putText(im0,person_id,(x1,y1-10),cv2.FONT_HERSHEY_COMPLEX,1,(0,0,255),2)                                unknow_person_emb[person_index]=this_frame_emb[index_]                                xyxy_all[person_index]=this_frame_xyxy[index_]                                this_person_index_no_state.append(person_index)                            # 行人匹配失敗 添加 特征和坐標(biāo)                            else:                                unknow_person_emb.append(this_frame_emb[index_])                                xyxy_all.append(this_frame_xyxy[index_])                        # 行人匹配失敗 添加 特征和坐標(biāo)                        else:                            unknow_person_emb.append(this_frame_emb[index_])                            xyxy_all.append(this_frame_xyxy[index_])        index+=1        log = 'all person:%s this frame person:%s '%(len(unknow_person_emb),len(this_frame_emb))        cv2.putText(im0, log, (20, 20), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)        cv2.namedWindow('image',cv2.WINDOW_NORMAL)        cv2.imshow('image', im0)        out.write(im0)        if cv2.waitKey(1) == ord('q'):            break
    video_capture.release()    out.release()    cv2.destroyAllWindows()

if __name__ == '__main__':    model_path = 'weights\yolov5m.pt'    hand_model_path = 'weights\hand_pose.h5'    Cam_num = r'test_video\Running - 294.mp4'    image_size = 416    conf_thres = 0.4    iou_thres = 0.4    device = 'cpu'    half = False    with torch.no_grad():        run()