單行人跟蹤

from utils.datasets import *from utils.utils import *import base64import osimport  tensorflow as tfimport torchfrom models.create_model import Create_Modelimport heapqos.environ['CUDA_VISIBLE_DEVICES'] = "-1"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)    return dis,list(dis <= tolerance)
# 獲得最相似的行人def get_top1(centerxy_list, centerxy):    centerxy_arr = np.array(centerxy_list)    cenarr = np.array(centerxy)
    dis = list(np.sqrt(np.sum(np.square(centerxy_arr - cenarr), axis=-1)))    min_num_index_list = map(dis.index, heapq.nsmallest(1, dis))    return list(min_num_index_list)

# 開始運(yùn)行def run():    # 加載yolo檢測模型    device = torch_utils.select_device('cpu')    model = torch.load(model_path, map_location=device)['model']    model.to(device).eval()
    # 加載行人重識別模型    input_size = (215, 90, 3)    model_, pred_model = Create_Model(inpt=input_size, num_classes=1812)    model_.load_weights('weights\ep039-loss0.066.h5')
    names = model.names if hasattr(model, 'names') else model.modules.names

    # 加載攝像頭    video_capture = cv2.VideoCapture(Cam_num)
    # 保存視頻的參數(shù)    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 #楨數(shù)    unknow_person_emb =[]  #未知的行人特征    center_xy = []  #行人目標(biāo)的中心點(diǎn)坐標(biāo)    know_center_xy_list=[] #已知行人的中心點(diǎn)坐標(biāo)    know_person_emb=None   #已知行人的特征

# 讀取視頻流數(shù)據(jù)    while True:        ret, im0 = video_capture.read()        iimage = im0.copy()        # 數(shù)據(jù)標(biāo)準(zhǔn)化        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)
        # 目標(biāo)檢測計算        pred = model(img, augment=False)[0]
        # uint8 to fp16/32        if half:            pred = pred.float()        # 將目標(biāo)檢測的結(jié)果進(jìn)行NMS處理，去掉多余的框        pred = non_max_suppression(pred, conf_thres, iou_thres,                                   fast=True, classes=None, agnostic=False)        # 循環(huán)處理獲取到的目標(biāo)數(shù)據(jù)        for i, det in enumerate(pred):            # 當(dāng)圖片中存在目標(biāo)時            if det is not None and len(det):                # 獲得目標(biāo)數(shù)據(jù)                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()                id = 0                # 獲得目標(biāo)的 坐標(biāo)，準(zhǔn)確度以及類別                for *xyxy, conf, cls in det:                    # 這里我們只針對人這一類別進(jìn)行計算                    if names[int(cls)] == 'person':                        # 獲得左上角 以及右下角的坐標(biāo)                        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))
                        # ROI數(shù)據(jù)標(biāo)準(zhǔn)化，進(jìn)入行人重識別別模型                        image_1 = np.asarray(roi).astype(np.float64) / 255                        photo1 = np.expand_dims(image_1 ,0)                        output1 = pred_model.predict(photo1)                        # 獲得中心點(diǎn)坐標(biāo)                        centerx = x_min+(x_max-x_min)/2                        centery = y_min+(y_max-y_min)/2                        id+=1
                        # 如果當(dāng)前是第一幀，則提取已知圖片的特則，                        # 這一步可以放到前面進(jìn)行處理                        if index ==0:                            unknow_person_emb.append(output1)                            center_xy.append([centerx,centery])                            know_img = cv2.imread('person.jpg')                            know_img = cv2.resize(know_img,(90,215))                            know_img = np.asarray(know_img).astype(np.float64) / 255                            know_img = np.expand_dims(know_img, 0)
                            know_person_emb = pred_model.predict(know_img)
                        # 如果已經(jīng)不是第一幀了，就進(jìn)行識別                        else:                            # 獲得最為接近的行人特征以及坐標(biāo)                            # 這里使用上一幀以及當(dāng)前楨的坐標(biāo)進(jìn)行對比                            # 一般來說兩楨之間，行人的移動距離不會太大，我們用這個來輔助判斷                            min_index=get_top1(center_xy,[centerx,centery])
                            # 獲得特征距離矩陣                            dit, com_p = com_person(unknow_person_emb, know_person_emb, tolerance=0.8)
                            # 如果最相似的行人的索引與特征距離矩陣中最小值的索引相同                            # 并且這個特征距離小于閾值，就說明找到了目標(biāo)                            if min_index[0] == dit.argmin() and com_p[dit.argmin()] == True:                                # 將識別結(jié)果繪制到圖片中                                label = 'ID-%s' % (1)                                plot_one_box(xyxy, im0, label=label, color=(0,255,0), line_thickness=5)
                                # 更新坐標(biāo)以及特征                                center_xy[min_index[0]]=[centerx,centery]                                know_person_emb = output1
                                # 只保存10個坐標(biāo)數(shù)據(jù) 用于繪制行走路徑                                if len(know_center_xy_list) < 10:                                    know_center_xy_list.append([centerx, centery])                                else:                                    know_center_xy_list.remove(know_center_xy_list[0])                                    know_center_xy_list.append([centerx, centery])
                                # 繪制目標(biāo)的行走路徑                                for i in range(1, len(know_center_xy_list)):                                    x1, y1 = know_center_xy_list[i - 1]                                    x2, y2 = know_center_xy_list[i]                                    cv2.line(im0, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)                            # 未找到該人,進(jìn)行繪制                            else:                                unknow_person_emb.append(output1)                                center_xy.append([centerx, centery])                                label = 'ID-%s' % ('unknow')                                plot_one_box(xyxy, im0, label=label, color=(0,0,255), line_thickness=5)        # 將每一幀的圖片進(jìn)行保存        out.write(im0)        index+=1        cv2.namedWindow('image',cv2.WINDOW_NORMAL)        cv2.imshow('image', im0)        if cv2.waitKey(1) == ord('q'):            break
    video_capture.release()    out.release()    cv2.destroyAllWindows()

if __name__ == '__main__':    # 基本設(shè)置    model_path = 'weights\yolov5s.pt'    hand_model_path = 'weights\hand_pose.h5'    Cam_num = r'test_video\person_test.mp4'    image_size = 416    conf_thres = 0.4    iou_thres = 0.4    device = 'cpu'    half = False    # 運(yùn)行    with torch.no_grad():        run()