SiamMask----demo.py详解

1_demo.py

 

本文对我修改过的demo.py文件进行修改，以图分析出如何与检测算法对接。

1 导入部分：

# --------------------------------------------------------
# SiamMask
# Licensed under The MIT License
# Written by Qiang Wang (wangqiang2015 at ia.ac.cn)
# --------------------------------------------------------
import glob
from tools.test import *   // 无需句点表示法:直接使用类，无需句点表示法. 可以参考  https://blog.csdn.net/gaoyueace/article/details/78700745

//1_创建解析对象
parser = argparse.ArgumentParser(description='PyTorch Tracking Demo')  // 建立解析对象 descirption是起了个名字

//2_添加参数
// resume梗概, required必选，type字符型,metavar说明中的参数名称，对于必选参数默认就是参数名称
parser.add_argument('--resume', default='', type=str, required=True,
                    metavar='PATH',help='path to latest checkpoint (default: none)')

// config配置  dest解析后的参数名称，默认情况下，对于可选参数选取最长的名称，中划线转换为下划线.
parser.add_argument('--config', dest='config', default='config_davis.json',
                    help='hyper-parameter of SiamMask in json format')

//
parser.add_argument('--base_path', default='../../data/tennis', help='datasets')

//3_解析
args = parser.parse_args()

 

2 主函数部分：

if __name__ == '__main__':
    # Setup device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 选择GPU
    torch.backends.cudnn.benchmark = True  # 优化运行效率

    # Setup Model
    cfg = load_config(args)  # 函数负责将参数取出来
    from custom import Custom  # siammask里面有custom.py文件，此文件里面有Custom类
    siammask = Custom(anchors=cfg['anchors'])
    if args.resume:
        assert isfile(args.resume), '{} is not a valid file'.format(args.resume)
        siammask = load_pretrain(siammask, args.resume)

    siammask.eval().to(device)

    # Parse Image file
    img_files = sorted(glob.glob(join(args.base_path, '*.jp*')))
    ims = [cv2.imread(imf) for imf in img_files]

    ###Modify###
    VeryBig=999999999  # 用于将视频框调整到最大
    Cap = cv2.VideoCapture(0)  # 设置读取摄像头
    # Cap.set(cv2.CAP_PROP_FRAME_WIDTH,VeryBig)
    # Cap.set(cv2.CAP_PROP_FRAME_HEIGHT,VeryBig)
    ret, frame = Cap.read()  # 读取帧
    ims = [frame] # 把frame放入列表格式的frame， 因为原文是将每帧图片放入列表
    ###Modify###

    # Select ROI
    cv2.namedWindow("SiamMask", cv2.WND_PROP_FULLSCREEN)
    # cv2.setWindowProperty("SiamMask", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
    try:
        init_rect = cv2.selectROI('SiamMask', ims[0], False, False)
        x, y, w, h = init_rect
    except:
        exit()

    toc = 0
###Modify### 此部分被下面替换
    # for f, im in enumerate(ims):  # 列表中逐个提取图片
    #     tic = cv2.getTickCount()
    #     if f == 0:  # init
    #         target_pos = np.array([x + w / 2, y + h / 2])
    #         target_sz = np.array([w, h])
    #         state = siamese_init(im, target_pos, target_sz, siammask, cfg['hp'])  # init tracker
    #     elif f > 0:  # tracking
    #         state = siamese_track(state, im, mask_enable=True, refine_enable=True)  # track
    #         location = state['ploygon'].flatten()
    #         mask = state['mask'] > state['p'].seg_thr

    #         im[:, :, 2] = (mask > 0) * 255 + (mask == 0) * im[:, :, 2]
    #         cv2.polylines(im, [np.int0(location).reshape((-1, 1, 2))], True, (0, 255, 0), 3)
    #         cv2.imshow('SiamMask', im)
    #         key = cv2.waitKey(1)
    #         if key > 0:
    #             break

    #     toc += cv2.getTickCount() - tic
    
    ###Modify###
    im=frame
    f=0
    target_pos = np.array([x + w / 2, y + h / 2])
    target_sz = np.array([w, h])
    state = siamese_init(im, target_pos, target_sz, siammask, cfg['hp'])  # init tracker
    while(True):
        tic = cv2.getTickCount()
        ret, im = Cap.read()  # 逐个提取frame
        if (ret==False):
            break;
        state = siamese_track(state, im, mask_enable=True, refine_enable=True)  # track
        location = state['ploygon'].flatten()
        mask = state['mask'] > state['p'].seg_thr

        im[:, :, 2] = (mask > 0) * 255 + (mask == 0) * im[:, :, 2]
        cv2.polylines(im, [np.int0(location).reshape((-1, 1, 2))], True, (0, 255, 0), 3)
        cv2.imshow('SiamMask', im)
        key = cv2.waitKey(1)
        if key > 0:
            break

        toc += cv2.getTickCount() - tic
        f=f+1
    ###Modify###
    
        

        
    toc /= cv2.getTickFrequency()
    fps = f / toc
    print('SiamMask Time: {:02.1f}s Speed: {:3.1f}fps (with visulization!)'.format(toc, fps))

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

1