手撕代码deep image matting(8):test

前面用了一堆的篇幅就是为了对dim训练过程的解读,但是训练了这么多次,再丑的媳妇也得见公婆,是骡子是马不测试一下谁也不知道这模型到底质量咋样。所以就该进入测试环节看看对于测试这块代码是怎么做的。

既然要找测试相关的东西,还是得先看一下说明文档。

那就有必要看看test.py里面的东西了。一如既往,把函数细节收起来,看看架构。

看起来有点像阉割后的训练步骤,毕竟测试集比训练集少了很多,只有1000个图片。那么他从config导入的路径参数都来看看有哪些。

 

 按照训练时候的路径分析,这几个路径分别对应的数据类型就极为明了:

fg_path_test: 测试集的前景图片

a_path_test:测试集的蒙版图片

bg_path_test:测试集的背景图片

out_path_test:前景和背景合成后的图片

剩下那些函数就在具体应用场景再看功能,直接看main函数

if __name__ == '__main__':
    checkpoint = 'BEST_checkpoint.tar'
    checkpoint = torch.load(checkpoint)
    model = checkpoint['model'].module
    model = model.to(device)
    model.eval()

    transformer = data_transforms['valid']

    names = gen_test_names()

    mse_losses = AverageMeter()
    sad_losses = AverageMeter()

    logger = get_logger()
    i = 0
    for name in tqdm(names):
        fcount = int(name.split('.')[0].split('_')[0])
        bcount = int(name.split('.')[0].split('_')[1])
        im_name = fg_test_files[fcount]
        # print(im_name)
        bg_name = bg_test_files[bcount]
        trimap_name = im_name.split('.')[0] + '_' + str(i) + '.png'
        # print('trimap_name: ' + str(trimap_name))

        trimap = cv.imread('data/Combined_Dataset/Test_set/Adobe-licensed images/trimaps/' + trimap_name, 0)
        # print('trimap: ' + str(trimap))

        i += 1
        if i == 20:
            i = 0

        img, alpha, fg, bg, new_trimap = process_test(im_name, bg_name, trimap)
        h, w = img.shape[:2]
        # mytrimap = gen_trimap(alpha)
        # cv.imwrite('images/test/new_im/'+trimap_name,mytrimap)

        x = torch.zeros((1, 4, h, w), dtype=torch.float)
        img = img[..., ::-1]  # RGB
        img = transforms.ToPILImage()(img)  # [3, 320, 320]
        img = transformer(img)  # [3, 320, 320]
        x[0:, 0:3, :, :] = img
        x[0:, 3, :, :] = torch.from_numpy(new_trimap.copy() / 255.)

        # Move to GPU, if available
        x = x.type(torch.FloatTensor).to(device)  # [1, 4, 320, 320]
        alpha = alpha / 255.

        with torch.no_grad():
            pred = model(x)  # [1, 4, 320, 320]

        pred = pred.cpu().numpy()
        pred = pred.reshape((h, w))  # [320, 320]

        pred[new_trimap == 0] = 0.0
        pred[new_trimap == 255] = 1.0
        cv.imwrite('images/test/out/' + trimap_name, pred * 255)

        # Calculate loss
        # loss = criterion(alpha_out, alpha_label)
        mse_loss = compute_mse(pred, alpha, trimap)
        sad_loss = compute_sad(pred, alpha)

        # Keep track of metrics
        mse_losses.update(mse_loss.item())
        sad_losses.update(sad_loss.item())
        print("sad:{} mse:{}".format(sad_loss.item(), mse_loss.item()))
        print("sad:{} mse:{}".format(sad_losses.avg, mse_losses.avg))
    print("sad:{} mse:{}".format(sad_losses.avg, mse_losses.avg))

加载了检查点的模型之后,把模型导入gpu,然后用transformer保存transforms,这里采用向data_transforms传入valid的处理方式进行。这时候的数据处理就只有两行:

 'valid': transforms.Compose([
        transforms.ToTensor(),#变成tensor
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])#标准化
    ]),

 在这之后 names = gen_test_names() 那就要看看gen_test_names这个函数是做什么的

def gen_test_names():
    num_fgs = 50
    num_bgs = 1000
    num_bgs_per_fg = 20

    names = []
    bcount = 0
    for fcount in range(num_fgs):
        for i in range(num_bgs_per_fg):
            names.append(str(fcount) + '_' + str(bcount) + '.png')
            bcount += 1

    return names

 这三个变量的数字基本上就能和整个测试集的图片数量对上:前景图片50个,背景图片1000个,那么他俩的比值就是1比20,也就是说1个前景合成了20个背景的数量最后产生1000个合成图片。那么这个函数其实就是出图片名字的列表。试验一下看看长啥样子。

 这个东西和之前的训练的时候train_names长得如出一辙,但是不一样的地方在于:训练数据集对于合成图片的命名规则的的确确是 前景图片位置_背景图片位置 这么定的,但是测试集的图片没有搞那么多,具体原因怎样只能往后来找。

关于logger日志这里今天先不管,直接看后面。mse_losses = AverageMeter()
    sad_losses = AverageMeter() 把mse_losses和sad_losses初始化之后,进入循环。

 for name in tqdm(names):
        fcount = int(name.split('.')[0].split('_')[0])
        bcount = int(name.split('.')[0].split('_')[1])
        im_name = fg_test_files[fcount]
        # print(im_name)
        bg_name = bg_test_files[bcount]
        trimap_name = im_name.split('.')[0] + '_' + str(i) + '.png'
        # print('trimap_name: ' + str(trimap_name))

        trimap = cv.imread('data/Combined_Dataset/Test_set/Adobe-licensed images/trimaps/' + trimap_name, 0)
        # print('trimap: ' + str(trimap))

        i += 1
        if i == 20:
            i = 0

        img, alpha, fg, bg, new_trimap = process_test(im_name, bg_name, trimap)
        h, w = img.shape[:2]
        # mytrimap = gen_trimap(alpha)
        # cv.imwrite('images/test/new_im/'+trimap_name,mytrimap)

        x = torch.zeros((1, 4, h, w), dtype=torch.float)
        img = img[..., ::-1]  # RGB
        img = transforms.ToPILImage()(img)  # [3, 320, 320]
        img = transformer(img)  # [3, 320, 320]
        x[0:, 0:3, :, :] = img
        x[0:, 3, :, :] = torch.from_numpy(new_trimap.copy() / 255.)

        # Move to GPU, if available
        x = x.type(torch.FloatTensor).to(device)  # [1, 4, 320, 320]
        alpha = alpha / 255.

        with torch.no_grad():
            pred = model(x)  # [1, 4, 320, 320]

        pred = pred.cpu().numpy()
        pred = pred.reshape((h, w))  # [320, 320]

        pred[new_trimap == 0] = 0.0
        pred[new_trimap == 255] = 1.0
        cv.imwrite('images/test/out/' + trimap_name, pred * 255)

        # Calculate loss
        # loss = criterion(alpha_out, alpha_label)
        mse_loss = compute_mse(pred, alpha, trimap)
        sad_loss = compute_sad(pred, alpha)

        # Keep track of metrics
        mse_losses.update(mse_loss.item())
        sad_losses.update(sad_loss.item())
        print("sad:{} mse:{}".format(sad_loss.item(), mse_loss.item()))
        print("sad:{} mse:{}".format(sad_losses.avg, mse_losses.avg))

fcount保存前景的编号,bcount保存背景的编号,这两个跟训练的时候思路是一样的,然后就是读取fg_test_files以及bg_test_files。这两个东西在data_gen里面被如此定义:

with open('Combined_Dataset/Test_set/test_fg_names.txt') as f:
    fg_test_files = f.read().splitlines()
with open('Combined_Dataset/Test_set/test_bg_names.txt') as f:
    bg_test_files = f.read().splitlines()

那就找到对应的部分看看具体是长什么样子。

 这个是test_bg_names.txt的内容,正好能对上测试集的背景图片名字。那么同理,fg_test_files也能对的上。那么按照训练数据集的处理思路,im_name = fg_test_files[fcount]  bg_name = bg_test_files[bcount]  这两句就能对上前景和背景的名字。那么 trimap_name = im_name.split('.')[0] + '_' + str(i) + '.png' 这句能出来的是什么内容,试试看。

 到了这里之后的下一句:trimap = cv.imread('data/Combined_Dataset/Test_set/Adobe-licensed images/trimaps/' + trimap_name, 0) 这才是重点。看到对应路径就寻找相应图片。

 原本的图片长这样:

 正好能对上相应的trimap图片,那么这里面的命名规则就是:

trimap_name = 前景图片的文件名_此时循环对应的参数i png

后续代码里面有一条:当i此时此刻因为加完了1之后如果等于20就进行重置为0 ,这是为什么?之前就提到每个前景图片都会和20个背景图片进行重组,那么一旦到了第21个就会换下一个前景图片进行合成。由于数组都是从0开始计算第一个,那么当数组下标到19时,此时的背景图片就会和新的前景图片合成。那么读取trimap的时候为了方便遍历读取就采取每个图片单独使用20个,编号从0到19,事实也的确如此。

 这里面观察一下提供好的merged_test文件夹,背景图会一直变,但是前景图合成了20个后就换下一个。那么img, alpha, fg, bg, new_trimap = process_test(im_name, bg_name, trimap) 就要看看process_test函数。

def process_test(im_name, bg_name, trimap):#im_name:前景文件名 bg_name:背景文件名 trimap:前景的trimap
    # print(bg_path_test + bg_name)
    im = cv.imread(fg_path_test + im_name)#读取前景图片
    a = cv.imread(a_path_test + im_name, 0)#读取蒙版图片
    h, w = im.shape[:2]#获得前景图片高和宽
    bg = cv.imread(bg_path_test + bg_name)#读取背景图片
    bh, bw = bg.shape[:2]#获取背景图片高和宽
    wratio = w / bw
    hratio = h / bh
    ratio = wratio if wratio > hratio else hratio
    if ratio > 1:
        bg = cv.resize(src=bg, dsize=(math.ceil(bw * ratio), math.ceil(bh * ratio)), interpolation=cv.INTER_CUBIC)


    return composite4_test(im, bg, a, w, h, trimap)

这里和前面的文章讲训练模型时候数据集处理的代码里面基本上一模一样,就是为了把背景图片改变成一个特别正好的尺寸包住前景。手撕代码deep image matting(6):dataset(1)_宇智波瞎眼猫的博客-CSDN博客

最后返回的东西就得仔细看看是怎么回事:

def composite4_test(fg, bg, a, w, h, trimap):
    fg = np.array(fg, np.float32)
    bg_h, bg_w = bg.shape[:2]
    x = max(0, int((bg_w - w) / 2))
    y = max(0, int((bg_h - h) / 2))
    crop = np.array(bg[y:y + h, x:x + w], np.float32)
    alpha = np.zeros((h, w, 1), np.float32)
    alpha[:, :, 0] = a / 255.
    # trimaps = np.zeros((h, w, 1), np.float32)
    # trimaps[:,:,0]=trimap/255.

    im = alpha * fg + (1 - alpha) * crop
    im = im.astype(np.uint8)

    new_a = np.zeros((bg_h, bg_w), np.uint8)
    new_a[y:y + h, x:x + w] = a
    new_trimap = np.zeros((bg_h, bg_w), np.uint8)
    new_trimap[y:y + h, x:x + w] = trimap
    cv.imwrite('images/test/new/' + trimap_name, new_trimap)
    new_im = bg.copy()
    new_im[y:y + h, x:x + w] = im
    # cv.imwrite('images/test/new_im/'+trimap_name,new_im)
    return new_im, new_a, fg, bg, new_trimap

虽然和训练数据集的composite4长得很像,但是有几个地方不太一样。在这两句x和y的设定上面,训练集选定的是随机值,这里面换成了更大值的选定。

 x = max(0, int((bg_w - w) / 2))
 y = max(0, int((bg_h - h) / 2))

在我们做前面的包住图片的处理的时候会有几种情况:

第一大类:前景图片比背景图片大

        1.1 w与bg_w的比值要比 h 和bg_h比值更大,此时就优先让bg_w与w持平,此时bg_h自然就能包住h也就是比h更大

        1.2 w与bg_w比值要比h与bg_h比值更小,此时优先让bg_h与h持平,此时bg_w肯定会比w更大

第二大类:前景图片比背景图片小,此时不需要对背景进行任何操作

在这几种情况之下,bg_w-w ,bg_h-h 以及0 的关系就有如下几种情况:

第一大类:前景图片比背景图片大

        1.1 w与bg_w的比值要比 h 和bg_h比值更大,此时 bg_w - w > 0 ,bg_h - h = 0

        1.2 w与bg_w比值要比h与bg_h比值更小,此时优先让bg_h与h持平,此时bg_w肯定会比w更大
                此时 bg_w - w = 0, bg_h - h > 0

第二大类:前景图片比背景图片小 此时 bg_w - w >= 0  bg_h - h >= 0

到这里分析一大堆其实x和y的起始位置确定的点其实很好说:必然在左上角某个比较细小的范围之内其实这里面切片背景的操作本质上和训练集的大差不差,只是缺少了随机性这一指标。然后就在这个起始点把 h w 尺寸的图片切出来,和蒙版组合到一起就是个新的合成图片,也就是 im 。

之后的操作稍微比之前要复杂了一些,不过本质上差不太多。直接上注释。

    new_a = np.zeros((bg_h, bg_w), np.uint8) #以背景的尺寸新建一个都为0的图片数组new_a
    new_a[y:y + h, x:x + w] = a # 对切片出来的区域赋值前景的蒙版值,因为此时新建的图片就是以该位置为起点进行合成
    new_trimap = np.zeros((bg_h, bg_w), np.uint8) 
    new_trimap[y:y + h, x:x + w] = trimap #对切片出来的位置赋值trimap

这些都完成之后, cv.imwrite('images/test/new/' + trimap_name, new_trimap) 在该位置把新的trimap生成出来。具体作用有多少,只能往后看看了。在这之后   new_im = bg.copy()
    new_im[y:y + h, x:x + w] = im 这两句把新的前景和背景的合成生成出来,尺寸还是背景的尺寸,但是图片是合成后的。这里面的处理和训练集就大相径庭:训练集是切出来的尺寸输入模型里面,测试集是要把背景尺寸的完整图片扔进去。

到了最后返回的这几个个参数:

new_im:新的合成图片,尺寸和背景图一致
new_a:新的和背景图尺寸一致的蒙版图片
fg:前景图片
bg:背景图片
new_trimap:新的trimap,尺寸和背景一致,精度没变

到这里整个process函数就搞定了。回到原来测试的代码里面:

        img, alpha, fg, bg, new_trimap = process_test(im_name, bg_name, trimap)
        h, w = img.shape[:2]
        # mytrimap = gen_trimap(alpha)
        # cv.imwrite('images/test/new_im/'+trimap_name,mytrimap)

        x = torch.zeros((1, 4, h, w), dtype=torch.float)
        img = img[..., ::-1]  # RGB
        img = transforms.ToPILImage()(img)  # [3, 320, 320]
        img = transformer(img)  # [3, 320, 320]
        x[0:, 0:3, :, :] = img
        x[0:, 3, :, :] = torch.from_numpy(new_trimap.copy() / 255.)

        # Move to GPU, if available
        x = x.type(torch.FloatTensor).to(device)  # [1, 4, 320, 320]
        alpha = alpha / 255.

这里面就要对应一下具体参数的:

img:对应new_im,新的合成图片,尺寸和背景图一致
alpha:对应new_a新的和背景图尺寸一致的蒙版图片
fg:前景图片
bg:背景图片
new_trimap:新的trimap,尺寸和背景一致,精度没变

h和w获得了新的img尺寸之后,初始化一个(h,w)尺寸的4通道矩阵x,再把img进行RGB重组,转化成PIL图片,然后就调用transformer也就是使用valid参数的data_transform。这里有一个地方和训练的不一样,x在初始化的时候采用的是 [1,4,h,w] 的尺寸,训练的时候没有前面的维度1。这里感觉需要再往后看的时候注意一下。官方在这块给的注释也有点奇怪,因为经过前面一系列的操作之后,图片只是按照前景图片的尺寸做了最佳适配,并没有训练集进行随机裁剪然后重组成320,320的尺寸。

x的第一维度按住,第二维度的前3个通道传入img数据,第4个通道传入新的归一化的trimap数据。然后将x传入gpu,蒙版值alpha归一化。然后就是开始进行计算的操作。

        with torch.no_grad():
            pred = model(x)  # [1, 4, 320, 320]

        pred = pred.cpu().numpy()
        pred = pred.reshape((h, w))  # [320, 320]

        pred[new_trimap == 0] = 0.0
        pred[new_trimap == 255] = 1.0
        cv.imwrite('images/test/out/' + trimap_name, pred * 255)

        # Calculate loss
        # loss = criterion(alpha_out, alpha_label)
        mse_loss = compute_mse(pred, alpha, trimap)
        sad_loss = compute_sad(pred, alpha)

        # Keep track of metrics
        mse_losses.update(mse_loss.item())
        sad_losses.update(sad_loss.item())
        print("sad:{} mse:{}".format(sad_loss.item(), mse_loss.item()))
        print("sad:{} mse:{}".format(sad_losses.avg, mse_losses.avg))

with torch.no_grad 我查到的作用就是不需要计算gradient进而节省gpu的计算空间,在这个情况下将x传入model走一遍模型。由于本身模型之前我们看训练DIMModel的时候,dimdataset返回的img规格是(4,320,320),但是这里面传入的数据是(1,4,h,w),这里面这个细节性的问题总感觉需要找个计算资源进行一波试验才能搞清楚。

经过model传入之后,最后出来的数据肯定是(h,w)尺寸的,因为4通道进去之后经过一系列的卷积出来就变成了1通道也就是(1,h,w),然后再经过squeeze把1通道的东西去掉就成了(h,w)尺寸。此时的pred是一个tensor,再将其转换成numpy并经历reshape,就成了(h,w)尺寸的数据。中间的写入图片的操作由于没有计算资源现在无法实验,所以直接看后面计算mse和sad

# compute the MSE error given a prediction, a ground truth and a trimap.
# pred: the predicted alpha matte
# target: the ground truth alpha matte
# trimap: the given trimap
#
def compute_mse(pred, alpha, trimap):
    num_pixels = float((trimap == 128).sum())
    return ((pred - alpha) ** 2).sum() / num_pixels


# compute the SAD error given a prediction and a ground truth.
#
def compute_sad(pred, alpha):
    diff = np.abs(pred - alpha)
    return np.sum(diff) / 1000

这里面有必要讲讲这两个关键指标:mse和sad。这一系列参数推荐这个文章来看看:

视频编码失真测度:SAD、SATD、SSD、MSE、PSNR_BigDream123的博客-CSDN博客

 mse是均方误差,sad为绝对误差和。把这两个算完之后,更新mse_losses以及sad_losses,再打印出来。到了最后,把mse和sad的平均值打印出来,整个测试就结束了。

后续需要填的坑有如下:

1. 为什么测试的时候数据集传入是四个维度。

2. 经过模型后出来的数据到底是什么样子。

3. 模型结束后写出来的trimap图片是什么样。

关于第三条在论文里面有成果展示:

 这个成果用最官方最有素质也是最书面的话表达了一个只要是发表论文的人都会在论文里面大书特书的观点:

老子牛逼。

评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值