【实例分割｜AdaptIS】数据集制作｜魔改 backbone

相关博客笔记

• 论文笔记：AdaptIS
• AdaptIS: Adaptive Instance Selection Network论文解读

数据集制作

'''
Auther: zth
Date: 2022-08-28 10:44:06
LastEditTime: 2022-08-28 11:21:07
Description: 从 labelme 标注的实例分割 json 标签文件生成适用于 AdaptIS 的数据集格式
'''
import json
import os
import numpy as np
import PIL.Image
from labelme import utils


def json_to_label(json_path):
    data = json.load(open(json_path))
    imageData = data['imageData']
    img = utils.img_b64_to_arr(imageData)
    label_name_to_value = {'_background_': 0, '1': 1}
    lbl, ins = utils.shapes_to_label(img.shape, data['shapes'],
                                     label_name_to_value)
    img_rgb = PIL.Image.fromarray(img)
    ins_pil = PIL.Image.fromarray(ins.astype(np.uint32), mode='I')

    return img_rgb, lbl, ins_pil


def main():
    jsons_path = "/Users/zth/Projects/0.线虫项目/2.制作COCO数据集/before"
    rgb_save_path = ""
    lbl_save_path = ""
    ins_pil_save_path = ""

    for dir, sub_dir, files in os.walk(jsons_path):
        if len(files) != 0:
            for file in files:
                img_rgb, lbl, ins_pil = json_to_label(file)
                name = os.path.splitext(file)[0]
                img_rgb_path = os.path.join(rgb_save_path, name + "_rgb.jpg")
                lbl_path = os.path.join(lbl_save_path, name + ".png")
                ins_pil_path = os.path.join(ins_pil_save_path,
                                            name + "_im.png")
                img_rgb.save(img_rgb_path)  # 保存 JPEG RGB 图像
                utils.lblsave(lbl_path, lbl)  # 保存语义分割掩码
                ins_pil.save(ins_pil_path)  # 保存实例标签


if __name__ == "__main__":
    main()

AdaIN

url: https://zhuanlan.zhihu.com/p/158657861
title: "AdaIN 笔记"
description: "论文 Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization的阅读笔记ICCV 2017的一篇论文，有点老，不过是一篇很棒的论文，做了非常多的实验，靠谱、实在，对我的研究也有非常大的帮助。…"
host: zhuanlan.zhihu.com
image: https://picx.zhimg.com/v2-7e1b8b9a5eeb9f8218f8f024dbe377ec_720w.jpg?source=172ae18b

相关 Github 项目

Fork

• MyAdaptis
• AdaptIS.CustomDataset

改进

• RoboTec
• iterdet

魔改 backbone

替换为 Unet 的 Pytorch 官方实现

• 官方的 Pytorch 实现中使用的是 Unet 论文中的官方实现，其中的跳连结构存在 crop 的操作，会导致部分像素的损失，所以改为 Pytorch 官方实现的 Unet。
• 步骤：
  • 步骤1：导入 Unet 网络文件 adaptis/model/toy/unet_pytorch.py
  • 步骤2：修改 adaptis/model/toy/models.py 文件，保证 input_channel=3，output_channel=32 即可
  • 步骤3：修改 train.py 中的 init_model() 函数

代码实现

步骤1

'''
Auther: zth
Date: 2022-09-24 14:01:10
LastEditTime: 2022-09-24 15:18:23
Description: 
'''
from typing import Dict
import torch
import torch.nn as nn
import torch.nn.functional as F


class DoubleConv(nn.Sequential):
    def __init__(self, in_channels, out_channels, mid_channels=None):
        if mid_channels is None:
            mid_channels = out_channels
        super(DoubleConv, self).__init__(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )


class Down(nn.Sequential):
    def __init__(self, in_channels, out_channels):
        super(Down, self).__init__(
            nn.MaxPool2d(2, stride=2),
            DoubleConv(in_channels, out_channels)
        )


class Up(nn.Module):
    def __init__(self, in_channels, out_channels, bilinear=True):
        super(Up, self).__init__()
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        x1 = self.up(x1)
        # [N, C, H, W]
        diff_y = x2.size()[2] - x1.size()[2]
        diff_x = x2.size()[3] - x1.size()[3]

        # padding_left, padding_right, padding_top, padding_bottom
        x1 = F.pad(x1, [diff_x // 2, diff_x - diff_x // 2,
                        diff_y // 2, diff_y - diff_y // 2])

        x = torch.cat([x2, x1], dim=1)
        x = self.conv(x)
        return x


class OutConv(nn.Sequential):
    def __init__(self, in_channels, num_classes):
        super(OutConv, self).__init__(
            nn.Conv2d(in_channels, num_classes, kernel_size=1)
        )


class UNetP(nn.Module):
    def __init__(self,
                 in_channels: int = 1,
                 num_classes: int = 2,
                 bilinear: bool = True,
                 base_c: int = 64):
        super(UNetP, self).__init__()
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.bilinear = bilinear
        self.feature_channels = base_c

        self.in_conv = DoubleConv(in_channels, base_c)
        self.down1 = Down(base_c, base_c * 2)
        self.down2 = Down(base_c * 2, base_c * 4)
        self.down3 = Down(base_c * 4, base_c * 8)
        factor = 2 if bilinear else 1
        self.down4 = Down(base_c * 8, base_c * 16 // factor)
        self.up1 = Up(base_c * 16, base_c * 8 // factor, bilinear)
        self.up2 = Up(base_c * 8, base_c * 4 // factor, bilinear)
        self.up3 = Up(base_c * 4, base_c * 2 // factor, bilinear)
        self.up4 = Up(base_c * 2, base_c, bilinear)
        self.out_conv = OutConv(base_c, num_classes)

    def get_feature_channels(self):
        return self.feature_channels

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        x1 = self.in_conv(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.out_conv(x)

        # return {"out": logits}
        return logits

步骤2

from adaptis.model.toy.unet_pytorch import UNetP

def get_unetp_model(channel_width=32, max_width=512, with_proposals=False, rescale_output=(0.2, -1.7), norm_layer=nn.BatchNorm2d):
    unet = UNetP(in_channels=3, num_classes=channel_width, bilinear=True, base_c=channel_width)
    in_channels = unet.get_feature_channels()

    return AdaptIS(
        backbone=unet,
        adaptis_head=ToyAdaptISHead(
            basic_blocks.SimpleConvController(3, in_channels, channel_width, norm_layer=norm_layer),
            in_channels,
            channels=channel_width,
            norm_radius=42,
            with_coord_features=True,
            rescale_output=rescale_output,
            norm_layer=norm_layer
        ),
        segmentation_head=basic_blocks.ConvHead(2, in_channels=in_channels, num_layers=3, norm_layer=norm_layer),
        proposal_head=basic_blocks.ConvHead(1, in_channels=in_channels, num_layers=2, norm_layer=norm_layer),
        with_proposals=with_proposals
    )

步骤3

def init_model():
    model_cfg = edict()
    model_cfg.syncbn = True

    model_cfg.input_normalization = {
        'mean': [0.5, 0.5, 0.5],
        'std': [0.5, 0.5, 0.5]
    }

    model_cfg.input_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(model_cfg.input_normalization['mean'],
                             model_cfg.input_normalization['std']),
    ])

    # training using DataParallel is not implemented
    norm_layer = torch.nn.BatchNorm2d

    # model = get_unet_model(norm_layer=norm_layer)
    model = get_unetp_model(norm_layer=norm_layer)
    model.apply(initializer.XavierGluon(rnd_type='gaussian', magnitude=1.0))

    return model, model_cfg