视频帧检测个人code记录

最新推荐文章于 2023-03-03 06:50:20 发布

CodeAIF

最新推荐文章于 2023-03-03 06:50:20 发布

阅读量752

点赞数

分类专栏：检测与追踪

本文链接：https://blog.csdn.net/qq_37553011/article/details/85122316

版权

检测与追踪专栏收录该内容

5 篇文章 0 订阅

订阅专栏

#检测模型调用函数： 基于detectorn 框架：
# 其中detection 是基于回归的框架
def prepare():
    if not torch.cuda.is_available():
        sys.exit("Need a CUDA device to run the code.")

    # cfg_file = 'voc2012.yaml'
    cfg_file = "FPN_voc2012.yaml"
    load_ckpt = "xx.pth"

    avi_dir = 'xx.path'


    dataset = datasets.get_minieye_dataset()
    print(cfg.MODEL.NUM_CLASSES)
    cfg.MODEL.NUM_CLASSES = len(dataset.classes)
    print('num_classes = ', cfg.MODEL.NUM_CLASSES)

    print('load cfg from file: ', cfg_file)
    cfg_from_file(cfg_file)

    cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False  # Don't need to load imagenet pretrained weights
    assert_and_infer_cfg()

    maskRCNN = Generalized_RCNN()

    maskRCNN.cuda()

    load_name = load_ckpt
    print("loading checkpoint %s" % (load_name))
    checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
    net_utils.load_ckpt(maskRCNN, checkpoint['model'])

    maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
                                 minibatch=True, device_ids=[0])  # only support single GPU
    return maskRCNN


if __name__ == '__main__':
    # imgpath = ''
    # img_path=input("imgpath:")
    #avi_dir = '  '
    avi_dir='-1/'
    save_img = 'DATA/vido_test/'
    colorsf= [(0, 255, 0), (0, 0, 255), (255, 0, 0)]
    AVI_path = glob.glob(avi_dir + '*.mp4')
    # print(AVI_path)
    vidos = filter(lambda x: x.endswith('mp4'), AVI_path)
    mask = prepare()
    for each_video in vidos:
        each_video_name = each_video.split('.')

        img_name = each_video_name[2].split('/')[1]
        #print(img_name)
        cap = cv2.VideoCapture(each_video)
        # print(cap
        frame_count = 4
        success = 1
        # vido=np.zeros((h,w,3),dtype='float16')
        c = 1
        rel = cap.isOpened()
        while rel:
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
            rel, frame = cap.read()
            currcentbox=detections(frame,mask,frame_count)
            twoframe=frame_count+1
            print('firstbox####', currcentbox)

            #############################################twoframe#############################
            cap.set(cv2.CAP_PROP_POS_FRAMES,twoframe)
            a,frame2=cap.read()
            twoframebox=detections(frame2,mask,twoframe)
            #newframe=frame.copy()
            if len(currcentbox)<len(twoframebox):
                # rectanglef(twoframebox,frame2,colorsf[1])
                # rectanglef(currcentbox,frame2,colorsf[0])
                lack_box=check_two_lack_box(currcentbox,twoframebox)
                print('-------------lack---',lack_box)
                print('---------------currcentbox----',currcentbox)
                for lbox in lack_box:
                    for cbox in currcentbox:
                        iou = calculateiou(lbox, cbox)
                        if iou>0.4:
                            currcentbox.remove(cbox)
                print('lack_box----------',lack_box)
                currcentbox=currcentbox+lack_box                                  #check_lack ---- two_frame  lack box  add frame box
                for box in currcentbox:
                    img=cv2.rectangle(frame, (box[1], box[2]), (box[3], box[4]), colorsf[0])
                for box in lack_box:
                    img=cv2.rectangle(frame, (box[1], box[2]), (box[3], box[4]), colorsf[2])      #  two :one lack box blue line
                cv2.imshow('frame-imshow-twoframe', img)
                cv2.waitKey(300)

            if len(currcentbox)>=len(twoframebox):
                more_current=check_two_lack_box(twoframebox,currcentbox)      # at currcentbox box look for twoframebox  more box
                print('more_current---------',more_current)
                for lbox in more_current:
                    for cbox in currcentbox:
                        iou = calculateiou(lbox, cbox)
                        if iou>0.4:
                            currcentbox.remove(cbox)
                print('-----more',more_current)
                for box in currcentbox:
                    imgs=cv2.rectangle(frame, (box[1], box[2]), (box[3], box[4]), colorsf[0])
                for box in more_current:
                    imgs=cv2.rectangle(frame,(box[1], box[2]), (box[3], box[4]), colorsf[1])              #two  more one lack box red line
                cv2.imshow('frame-imshow-twoframe', imgs)
                cv2.waitKey(300)
            frame_count += 2