PyTorch 新たなクラスの物体検出をSSDでやってみる

今回は、血液の顕微鏡画像から細胞を検出するSSDモデルを作ってみたいと思います。

こんにちは cedro です。

先回、SSDの学習済みモデルを使った物体検出を行ってみましたが、物体検出できるのはあらかじめ学習した20クラスだけです。

新たなクラスの物体検出をするには、どうしたら良いのでしょうか。大量のアノテーション付きデータを用意してゼロから学習するしかないのでしょうか。

そんなことはありません。SSDのネットワークの前半はVGG16の学習済みモデルの一部をベースネットワークとして使えるので、後半部分のみ新たに検出したいクラスのデータセットを少量学習させることによって、新しいSSDモデルを構築できます。

ということで、今回は、血液の顕微鏡画像から細胞を検出するSSDモデルを作ってみたいと思います。

データセットを準備します

今回、使用するデータセットは、BCCD Dataset という血液の顕微鏡写真で、白血球、赤血球、血小板の３つについてバウンディングボックスのアノテーションデータが付いたものです。

データセットの仕様が、PASCAL Visual Object Classes ですので、PyTorch のSSDモデルで簡単に読み込むことが出来ます。

物体検出をした時のイメージは、こんな感じ。wbc が白血球、rbc が赤血球、platelets が血小板です。それにしても、なんともマニアックなデータセットですよね。

BCCD Dataset はこういった画像とアノテーションデータのセットが、全部で364個（trainval:292個、test:72個）しかない非常に小さなデータセットです。

Github から BCCD Datset をダウンロードします。実際に使用するのは、赤枠で囲ったBCCDフォルダーの部分だけです。

コードを書きます

今回も、PyTorchニューラルネットワーク実装ハンドブックのお世話になります。Github からサンプルコードをクローンあるいはダウンロードします。今回使用するのは、Chapter7です。

chapter7フォルダーの中身です。今回は、この chapter7フォルダーの中で、コードの追加・修正を行います。

まず、VOCdevkitフォルダーを追加し、その中に先程ダウンロードした BCCD フォルダーを格納します。

そして、CNNのベースネットワーク（ vgg16_reducedfc.pth ）をダウンロードし、weights フォルダーに格納しておきます。

これが、SSDのネットワーク構成です。前半は、VGG16ネットワークの学習済みの重みの一部をそのまま利用し、後半のExtra Feature Layers のみを新たなデータセットを使って学習することで、新たなクラスの物体検出が出来るSSDネットワークが出来ます。早速、コードを書いてみましょう。

from data import *
from utils.augmentations import SSDAugmentation
from layers.modules import MultiBoxLoss
from ssd import build_ssd
import os
import sys
import time
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
import torch.nn.init as init
import torch.utils.data as data
import numpy as np
import argparse

args = {'dataset':'BCCD',  # VOC → BCCD
        'basenet':'vgg16_reducedfc.pth',
        'batch_size':12,
        'resume':'',
        'start_iter':0,
        'num_workers':0,  # 4 → 0
        'cuda':True,  # Macの場合False
        'lr':5e-4,
        'momentum':0.9,
        'weight_decay':5e-4,
        'gamma':0.1,
        'save_folder':'weights/'
       }

if torch.cuda.is_available():
    if args['cuda']:
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
    if not args['cuda']:
        print("WARNING: It looks like you have a CUDA device, but aren't " +
              "using CUDA.\nRun with --cuda for optimal training speed.")
        torch.set_default_tensor_type('torch.FloatTensor')
else:
    torch.set_default_tensor_type('torch.FloatTensor')

if not os.path.exists(args['save_folder']):
    os.mkdir(args['save_folder'])

# 訓練データの読み込み
cfg = voc
dataset = VOCDetection(root=VOC_ROOT,
                       transform=SSDAugmentation(cfg['min_dim'],
                                                 MEANS))

# ネットワークの定義
ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
net = ssd_net
print(net)

if args['cuda']:
    net = torch.nn.DataParallel(ssd_net)
    cudnn.benchmark = True
    
# パラメータのロード
if args['resume']:
    print('Resuming training, loading {}...'.format(args['resume']))
    ssd_net.load_weights(args['resume'])
else:
    vgg_weights = torch.load(args['save_folder'] + args['basenet'])
    print('Loading base network...')
    ssd_net.vgg.load_state_dict(vgg_weights)
    
if args['cuda']:
    net = net.cuda()
    
def adjust_learning_rate(optimizer, gamma, step):
    lr = args['lr'] * (gamma ** (step))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def xavier(param):
    init.xavier_uniform(param)

def weights_init(m):
    if isinstance(m, nn.Conv2d):
        xavier(m.weight.data)
        m.bias.data.zero_()
        
if not args['resume']:
    print('Initializing weights...')
    # initialize newly added layers' weights with xavier method
    ssd_net.extras.apply(weights_init)
    ssd_net.loc.apply(weights_init)
    ssd_net.conf.apply(weights_init)
    
# 最適化パラメータの設定
optimizer = optim.SGD(net.parameters(), lr=args['lr'], momentum=args['momentum'],
                      weight_decay=args['weight_decay'])

# 損失関数の設定
criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,
                         False, args['cuda'])

net.train()
# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
print('Loading the dataset...')

epoch_size = len(dataset) // args['batch_size']
print('Training SSD on:', dataset.name)
print('Using the specified args:')
print(args)

step_index = 0

# 訓練データの読み込み
data_loader = data.DataLoader(dataset, args['batch_size'],
                              num_workers=args['num_workers'],
                              shuffle=True, collate_fn=detection_collate,
                              pin_memory=True)

# 学習の開始
batch_iterator = None
# iterationでループして、cfg['max_iter']まで学習する
for iteration in range(args['start_iter'], cfg['max_iter']):
    # 学習開始時または1epoch終了後にdata_loaderから訓練データをロードする
    if (not batch_iterator) or (iteration % epoch_size ==0):
      batch_iterator = iter(data_loader)
      loc_loss = 0
      conf_loss = 0
      epoch += 1

    if iteration in cfg['lr_steps']:
        step_index += 1
        adjust_learning_rate(optimizer, args['gamma'], step_index)

    # load train data
    # バッチサイズ分の訓練データをload
    images, targets = next(batch_iterator)

    if args['cuda']:
        images = Variable(images.cuda())
        targets = [Variable(ann.cuda(), volatile=True) for ann in targets]
    else:
        images = Variable(images)
        targets = [Variable(ann, volatile=True) for ann in targets]
    # forward
    t0 = time.time()
    out = net(images)
    # backprop
    optimizer.zero_grad()
    loss_l, loss_c = criterion(out, targets)
    loss = loss_l + loss_c
    loss.backward()
    optimizer.step()
    t1 = time.time()
    loc_loss += loss_l.item()
    conf_loss += loss_c.item()

    #ログの出力
    if iteration % 10 == 0:
        print('timer: %.4f sec.' % (t1 - t0))
        print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.item()), end=' ')

# 学習済みモデルの保存
torch.save(ssd_net.state_dict(),
           args['save_folder'] + '' + args['dataset'] + '.pth')

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

from data import *

from utils.augmentations import SSDAugmentation

from layers.modules import MultiBoxLoss

from ssd import build_ssd

import os

import sys

import time

import torch

from torch.autograd import Variable

import torch.nn as nn

import torch.optim as optim

import torch.backends.cudnn as cudnn

import torch.nn.init as init

import torch.utils.data as data

import numpy as np

import argparse

args = {'dataset':'BCCD', # VOC → BCCD

'basenet':'vgg16_reducedfc.pth',

'batch_size':12,

'resume':'',

'start_iter':0,

'num_workers':0, # 4 → 0

'cuda':True, # Macの場合False

'lr':5e-4,

'momentum':0.9,

'weight_decay':5e-4,

'gamma':0.1,

'save_folder':'weights/'

}

if torch.cuda.is_available():

if args['cuda']:

torch.set_default_tensor_type('torch.cuda.FloatTensor')

if not args['cuda']:

print("WARNING: It looks like you have a CUDA device, but aren't " +

"using CUDA.\nRun with --cuda for optimal training speed.")

torch.set_default_tensor_type('torch.FloatTensor')

else:

torch.set_default_tensor_type('torch.FloatTensor')

if not os.path.exists(args['save_folder']):

os.mkdir(args['save_folder'])

# 訓練データの読み込み

cfg = voc

dataset = VOCDetection(root=VOC_ROOT,

transform=SSDAugmentation(cfg['min_dim'],

MEANS))

# ネットワークの定義

ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes'])

net = ssd_net

print(net)

if args['cuda']:

net = torch.nn.DataParallel(ssd_net)

cudnn.benchmark = True

# パラメータのロード

if args['resume']:

print('Resuming training, loading {}...'.format(args['resume']))

ssd_net.load_weights(args['resume'])

else:

vgg_weights = torch.load(args['save_folder'] + args['basenet'])

print('Loading base network...')

ssd_net.vgg.load_state_dict(vgg_weights)

if args['cuda']:

net = net.cuda()

def adjust_learning_rate(optimizer, gamma, step):

lr = args['lr'] * (gamma ** (step))

for param_group in optimizer.param_groups:

param_group['lr'] = lr

def xavier(param):

init.xavier_uniform(param)

def weights_init(m):

if isinstance(m, nn.Conv2d):

xavier(m.weight.data)

m.bias.data.zero_()

if not args['resume']:

print('Initializing weights...')

# initialize newly added layers' weights with xavier method

ssd_net.extras.apply(weights_init)

ssd_net.loc.apply(weights_init)

ssd_net.conf.apply(weights_init)

# 最適化パラメータの設定

optimizer = optim.SGD(net.parameters(), lr=args['lr'], momentum=args['momentum'],

weight_decay=args['weight_decay'])

# 損失関数の設定

criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,

False, args['cuda'])

net.train()

# loss counters

loc_loss = 0

conf_loss = 0

epoch = 0

print('Loading the dataset...')

epoch_size = len(dataset) // args['batch_size']

print('Training SSD on:', dataset.name)

print('Using the specified args:')

print(args)

step_index = 0

# 訓練データの読み込み

data_loader = data.DataLoader(dataset, args['batch_size'],

num_workers=args['num_workers'],

shuffle=True, collate_fn=detection_collate,

pin_memory=True)

# 学習の開始

batch_iterator = None

# iterationでループして、cfg['max_iter']まで学習する

for iteration in range(args['start_iter'], cfg['max_iter']):

# 学習開始時または1epoch終了後にdata_loaderから訓練データをロードする

if (not batch_iterator) or (iteration % epoch_size ==0):

batch_iterator = iter(data_loader)

loc_loss = 0

conf_loss = 0

epoch += 1

if iteration in cfg['lr_steps']:

step_index += 1

adjust_learning_rate(optimizer, args['gamma'], step_index)

# load train data

# バッチサイズ分の訓練データをload

images, targets = next(batch_iterator)

if args['cuda']:

images = Variable(images.cuda())

targets = [Variable(ann.cuda(), volatile=True) for ann in targets]

else:

images = Variable(images)

targets = [Variable(ann, volatile=True) for ann in targets]

# forward

t0 = time.time()

out = net(images)

# backprop

optimizer.zero_grad()

loss_l, loss_c = criterion(out, targets)

loss = loss_l + loss_c

loss.backward()

optimizer.step()

t1 = time.time()

loc_loss += loss_l.item()

conf_loss += loss_c.item()

#ログの出力

if iteration % 10 == 0:

print('timer: %.4f sec.' % (t1 - t0))

print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.item()), end=' ')

# 学習済みモデルの保存

torch.save(ssd_net.state_dict(),

args['save_folder'] + '' + args['dataset'] + '.pth')

学習するためのコードです。train.py という名前で Chapter7 に保存します。

今回は、Windowsで動かす想定をしています。もし、Macで動かす場合は、24行目を ‘cuda’ : False に変更して下さい。

from .config import HOME
import os.path as osp
import sys
import torch
import torch.utils.data as data
import cv2
import numpy as np
if sys.version_info[0] == 2:
    import xml.etree.cElementTree as ET
else:
    import xml.etree.ElementTree as ET

VOC_CLASSES = ('rbc', 'wbc', 'platelets')  # 修正

dir_cur = osp.dirname(__file__)
dir_voc = osp.join(dir_cur, "..", "VOCdevkit")
VOC_ROOT = osp.abspath(dir_voc)

class VOCAnnotationTransform(object):
    def __init__(self, class_to_ind=None, keep_difficult=False):
        self.class_to_ind = class_to_ind or dict(
            zip(VOC_CLASSES, range(len(VOC_CLASSES))))
        self.keep_difficult = keep_difficult

    def __call__(self, target, width, height):
        res = []
        for obj in target.iter('object'):
            difficult = int(obj.find('difficult').text) == 1
            if not self.keep_difficult and difficult:
                continue
            name = obj.find('name').text.lower().strip()
            bbox = obj.find('bndbox')

            pts = ['xmin', 'ymin', 'xmax', 'ymax']
            bndbox = []
            for i, pt in enumerate(pts):
                cur_pt = int(bbox.find(pt).text) - 1
                cur_pt = cur_pt / width if i % 2 == 0 else cur_pt / height
                bndbox.append(cur_pt)
            label_idx = self.class_to_ind[name]
            bndbox.append(label_idx)
            res += [bndbox]  
        return res  


class VOCDetection(data.Dataset):
    def __init__(self, root,
                 image_sets=[('BCCD', 'trainval')],  # 修正  
                 transform=None, target_transform=VOCAnnotationTransform(),
                 dataset_name='VOC0712'):
        self.root = root
        self.image_set = image_sets
        self.transform = transform
        self.target_transform = target_transform
        self.name = dataset_name
        self._annopath = osp.join('%s', 'Annotations', '%s.xml')
        self._imgpath = osp.join('%s', 'JPEGImages', '%s.jpg')
        self.ids = list()
        for (dir, name) in image_sets:  # 修正
            rootpath = osp.join(self.root, dir)  # 修正
            for line in open(osp.join(rootpath, 'ImageSets', 'Main', name + '.txt')):
                self.ids.append((rootpath, line.strip()))

    def __getitem__(self, index):
        im, gt, h, w = self.pull_item(index)

        return im, gt

    def __len__(self):
        return len(self.ids)

    def pull_item(self, index):
        img_id = self.ids[index]

        target = ET.parse(self._annopath % img_id).getroot()
        img = cv2.imread(self._imgpath % img_id)
        height, width, channels = img.shape

        if self.target_transform is not None:
            target = self.target_transform(target, width, height)

        if self.transform is not None:
            target = np.array(target)
            img, boxes, labels = self.transform(img, target[:, :4], target[:, 4])
            img = img[:, :, (2, 1, 0)]
            target = np.hstack((boxes, np.expand_dims(labels, axis=1)))
        return torch.from_numpy(img).permute(2, 0, 1), target, height, width

    def pull_image(self, index):
        img_id = self.ids[index]
        return cv2.imread(self._imgpath % img_id, cv2.IMREAD_COLOR)

    def pull_anno(self, index):
        img_id = self.ids[index]
        anno = ET.parse(self._annopath % img_id).getroot()
        gt = self.target_transform(anno, 1, 1)
        return img_id[1], gt

    def pull_tensor(self, index):
        return torch.Tensor(self.pull_image(index)).unsqueeze_(0)

100

101

from .config import HOME

import os.path as osp

import sys

import torch

import torch.utils.data as data

import cv2

import numpy as np

if sys.version_info[0] == 2:

import xml.etree.cElementTree as ET

else:

import xml.etree.ElementTree as ET

VOC_CLASSES = ('rbc', 'wbc', 'platelets') # 修正

dir_cur = osp.dirname(__file__)

dir_voc = osp.join(dir_cur, "..", "VOCdevkit")

VOC_ROOT = osp.abspath(dir_voc)

class VOCAnnotationTransform(object):

def __init__(self, class_to_ind=None, keep_difficult=False):

self.class_to_ind = class_to_ind or dict(

zip(VOC_CLASSES, range(len(VOC_CLASSES))))

self.keep_difficult = keep_difficult

def __call__(self, target, width, height):

res = []

for obj in target.iter('object'):

difficult = int(obj.find('difficult').text) == 1

if not self.keep_difficult and difficult:

continue

name = obj.find('name').text.lower().strip()

bbox = obj.find('bndbox')

pts = ['xmin', 'ymin', 'xmax', 'ymax']

bndbox = []

for i, pt in enumerate(pts):

cur_pt = int(bbox.find(pt).text) - 1

cur_pt = cur_pt / width if i % 2 == 0 else cur_pt / height

bndbox.append(cur_pt)

label_idx = self.class_to_ind[name]

bndbox.append(label_idx)

res += [bndbox]

return res

class VOCDetection(data.Dataset):

def __init__(self, root,

image_sets=[('BCCD', 'trainval')], # 修正

transform=None, target_transform=VOCAnnotationTransform(),

dataset_name='VOC0712'):

self.root = root

self.image_set = image_sets

self.transform = transform

self.target_transform = target_transform

self.name = dataset_name

self._annopath = osp.join('%s', 'Annotations', '%s.xml')

self._imgpath = osp.join('%s', 'JPEGImages', '%s.jpg')

self.ids = list()

for (dir, name) in image_sets: # 修正

rootpath = osp.join(self.root, dir) # 修正

for line in open(osp.join(rootpath, 'ImageSets', 'Main', name + '.txt')):

self.ids.append((rootpath, line.strip()))

def __getitem__(self, index):

im, gt, h, w = self.pull_item(index)

return im, gt

def __len__(self):

return len(self.ids)

def pull_item(self, index):

img_id = self.ids[index]

target = ET.parse(self._annopath % img_id).getroot()

img = cv2.imread(self._imgpath % img_id)

height, width, channels = img.shape

if self.target_transform is not None:

target = self.target_transform(target, width, height)

if self.transform is not None:

target = np.array(target)

img, boxes, labels = self.transform(img, target[:, :4], target[:, 4])

img = img[:, :, (2, 1, 0)]

target = np.hstack((boxes, np.expand_dims(labels, axis=1)))

return torch.from_numpy(img).permute(2, 0, 1), target, height, width

def pull_image(self, index):

img_id = self.ids[index]

return cv2.imread(self._imgpath % img_id, cv2.IMREAD_COLOR)

def pull_anno(self, index):

img_id = self.ids[index]

anno = ET.parse(self._annopath % img_id).getroot()

gt = self.target_transform(anno, 1, 1)

return img_id[1], gt

def pull_tensor(self, index):

return torch.Tensor(self.pull_image(index)).unsqueeze_(0)

dataフォルダーの中にある、voc0712.py のコードを一部修正します。13行目のVOC_CLASSES はBCCDのデータセットに合わせて修正します。

48行目 image_sets の指定を’BCCD’, ‘trainval’ のみにします。59行目の year は dir に変更し、60行目の ‘VOC’＋year はdir に変更します。

# SSD300 CONFIGS
voc = {
    'num_classes': 21,
    # handbook
    #'lr_steps': (80000, 100000, 120000),
    #'max_iter': 120000,
    'lr_steps': (8000, 10000, 12000),
    'max_iter': 3000,  # 12000 → 3000
    # handbook
    'feature_maps': [38, 19, 10, 5, 3, 1],
    'min_dim': 300,
    'steps': [8, 16, 32, 64, 100, 300],
    'min_sizes': [30, 60, 111, 162, 213, 264],
    'max_sizes': [60, 111, 162, 213, 264, 315],
    'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]],
    'variance': [0.1, 0.2],
    'clip': True,
    'name': 'BCCD',  # 'VOC' → 'BCCD' 
}

# SSD300 CONFIGS

voc = {

'num_classes': 21,

# handbook

#'lr_steps': (80000, 100000, 120000),

#'max_iter': 120000,

'lr_steps': (8000, 10000, 12000),

'max_iter': 3000, # 12000 → 3000

# handbook

'feature_maps': [38, 19, 10, 5, 3, 1],

'min_dim': 300,

'steps': [8, 16, 32, 64, 100, 300],

'min_sizes': [30, 60, 111, 162, 213, 264],

'max_sizes': [60, 111, 162, 213, 264, 315],

'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]],

'variance': [0.1, 0.2],

'clip': True,

'name': 'BCCD', # 'VOC' → 'BCCD'

}

dataフォルダー内にある、config.py の SSD300 CONFIGS の一部を修正します。8行目 ‘max_iter’ :12000 → 3000 に変更し、18行目 ‘VOC’ → ‘BCCD’ に変更します。

学習・推論を実行します

train.py を実行します。trainvalデータ292個、3000 epoch をGTX1060で学習を行ったところ、約1時間で完了しました。

学習が完了したら、weights フォルダーに保存される重みファイル（BCCD.pth）を使って、テスト画像の推論を実行します。

import os
import sys
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
import numpy as np
import cv2
if torch.cuda.is_available():
    torch.set_default_tensor_type('torch.cuda.FloatTensor')

from ssd import build_ssd

# SSDネットワークの定義と重みファイルのロード
net = build_ssd('test', 300, 21)    
net.load_weights('./weights/BCCD.pth')

from matplotlib import pyplot as plt
from data import VOCDetection, VOC_ROOT, VOCAnnotationTransform

# BCCD_test 読み込み
testset = VOCDetection(VOC_ROOT, [('BCCD', 'test')], None, VOCAnnotationTransform())
img_id = 25  
image = testset.pull_image(img_id)

# テスト画像の表示
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.figure(figsize=(10,10))
plt.imshow(rgb_image)
plt.show()

x = cv2.resize(image, (300, 300)).astype(np.float32)  # 300*300にリサイズ
x -= (104.0, 117.0, 123.0)
x = x.astype(np.float32)
x = x[:, :, ::-1].copy()
x = torch.from_numpy(x).permute(2, 0, 1)  # [300,300,3] → [３,300,300]
xx = Variable(x.unsqueeze(0))     # [3,300,300] → [1,3,300,300]
if torch.cuda.is_available():
    xx = xx.cuda()
# 順伝播を実行し、推論結果を出力    
y = net(xx)

from data import VOC_CLASSES as labels

plt.figure(figsize=(10,10))
colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
plt.imshow(rgb_image) 
currentAxis = plt.gca()
# 推論結果をdetectionsに格納
detections = y.data
# scale each detection back up to the image
scale = torch.Tensor(rgb_image.shape[1::-1]).repeat(2)

# バウンディングボックスとクラス名の表示
for i in range(detections.size(1)):
    j = 0
    # 確信度confが0.6以上のボックスを表示
    # jは確信度上位200件のボックスのインデックス
    # detections[0,i,j]は[conf,xmin,ymin,xmax,ymax]の形状
    while detections[0,i,j,0] >= 0.6:
        score = detections[0,i,j,0]
        label_name = labels[i-1]
        display_txt = '%s: %.2f'%(label_name, score)
        pt = (detections[0,i,j,1:]*scale).cpu().numpy()
        coords = (pt[0], pt[1]), pt[2]-pt[0]+1, pt[3]-pt[1]+1
        color = colors[i]
        currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))
        currentAxis.text(pt[0], pt[1], display_txt, bbox={'facecolor':color, 'alpha':0.5})
        j+=1
        
plt.show()
plt.close()

import os

import sys

import torch

import torch.nn as nn

import torch.backends.cudnn as cudnn

from torch.autograd import Variable

import numpy as np

import cv2

if torch.cuda.is_available():

torch.set_default_tensor_type('torch.cuda.FloatTensor')

from ssd import build_ssd

# SSDネットワークの定義と重みファイルのロード

net = build_ssd('test', 300, 21)

net.load_weights('./weights/BCCD.pth')

from matplotlib import pyplot as plt

from data import VOCDetection, VOC_ROOT, VOCAnnotationTransform

# BCCD_test 読み込み

testset = VOCDetection(VOC_ROOT, [('BCCD', 'test')], None, VOCAnnotationTransform())

img_id = 25

image = testset.pull_image(img_id)

# テスト画像の表示

rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

plt.figure(figsize=(10,10))

plt.imshow(rgb_image)

plt.show()

x = cv2.resize(image, (300, 300)).astype(np.float32) # 300*300にリサイズ

x -= (104.0, 117.0, 123.0)

x = x.astype(np.float32)

x = x[:, :, ::-1].copy()

x = torch.from_numpy(x).permute(2, 0, 1) # [300,300,3] → [３,300,300]

xx = Variable(x.unsqueeze(0)) # [3,300,300] → [1,3,300,300]

if torch.cuda.is_available():

xx = xx.cuda()

# 順伝播を実行し、推論結果を出力

y = net(xx)

from data import VOC_CLASSES as labels

plt.figure(figsize=(10,10))

colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()

plt.imshow(rgb_image)

currentAxis = plt.gca()

# 推論結果をdetectionsに格納

detections = y.data

# scale each detection back up to the image

scale = torch.Tensor(rgb_image.shape[1::-1]).repeat(2)

# バウンディングボックスとクラス名の表示

for i in range(detections.size(1)):

j = 0

# 確信度confが0.6以上のボックスを表示

# jは確信度上位200件のボックスのインデックス

# detections[0,i,j]は[conf,xmin,ymin,xmax,ymax]の形状

while detections[0,i,j,0] >= 0.6:

score = detections[0,i,j,0]

label_name = labels[i-1]

display_txt = '%s: %.2f'%(label_name, score)

pt = (detections[0,i,j,1:]*scale).cpu().numpy()

coords = (pt[0], pt[1]), pt[2]-pt[0]+1, pt[3]-pt[1]+1

color = colors[i]

currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))

currentAxis.text(pt[0], pt[1], display_txt, bbox={'facecolor':color, 'alpha':0.5})

j+=1

plt.show()

plt.close()

推論を実行するコードです。inference.py という名前で、chapter7フォルダーに保存します。

26行目でBCCD_test を読み込んでいますので、27行目のimg_id を指定することで、物体検出に使うテストデータを選択できます。

それでは、inference.py を実行してみます。

テストデータの42番目の画像です。これを、物体検出させると、

若干取りこぼしがありそうですが、まあまあ物体検出できている感じです。

もう１つ行ってみましょう。テストデータの55番目の画像です。これを、物体検出させると、

これも、まずまずでしょうか。

たった、292個のデータを学習させただけですが、ベースネットワークは1000クラスの分類を学習したVGG16ネットワークなので、結構物体検出ができるものですね。

では、また。

google colab バージョンを追加

2021/1

　上記でご説明したコードをGoogle Colabで動かす形にしてGithubに上げました。この「リンク」をクリックし表示されたノートブックの先頭にある「Colab on Web」ボタンをクリックすると動かせます。BCCDデータセット、CNNのベースネットワークの重みもコードに保存済みなので、手軽に試せると思います。

突然の質問失礼します。

train.py

のコード1行目の
from data import　*
はどういう意味でしょうか？

下記の様なエラーが出ており、ご紹介していただいている学習を進めることができずに困っています。
importの先に何か、importするデータを指定しないといけないと思うのですが、初学者でして、自力で解決できず困っています。
よろしければご教授願えれば幸いです。

No module named ‘torch’
File “/Users/myname/Desktop/pytorch_handbook/chapter7/data/voc0712.py”, line 17, in
import torch
File “/Users/myname/Desktop/pytorch_handbook/chapter7/data/__init__.py”, line 7, in
from .voc0712 import VOCDetection, VOCAnnotationTransform, VOC_CLASSES, VOC_ROOT
File “/Users/myname/Desktop/pytorch_handbook/chapter7/train.py”, line 1, in
from data import *

PyTorch 新たなクラスの物体検出をSSDでやってみる

今回は、血液の顕微鏡画像から細胞を検出するSSDモデルを作ってみたいと思います。

データセットを準備します

コードを書きます

学習・推論を実行します

google colab バージョンを追加

2 件のコメント

コメントを残すコメントをキャンセル

ABOUTこの記事をかいた人

NEW POSTこのライターの最新記事

LightWeight_GANで、新しいラーメンを創作してみる

AIで画像から鉛筆画を描く

CPUだけで顔・手・ポーズのリアルタイム検出を行う

ディープラーニングで、1枚の顔画像から詳細3Dモデルを作って自由に動かす

最近の投稿

最近のコメント

アーカイブ

カテゴリー

メタ情報

今回は、血液の顕微鏡画像から細胞を検出するSSDモデルを作ってみたいと思います。

データセットを準備します

コードを書きます

学習・推論を実行します

google colab バージョンを追加

2 件のコメント

コメントを残す コメントをキャンセル

RECOMMENDこちらの記事も人気です。

たった100枚の画像でもGANの学習を可能にする技術

Keras MLPを改造して定番パターンを勉強する

学習済みWord2Vec モデルをサクッと使ってみる

Keras KMNIST「くずし文字」練習器を作ってみる

PULSEで低解像度の顔画像を高解像度に変換する

SONY Neural Network Console で半教師あり学習…

PyTorch そして LSGANをやってみる

Keras LSTM の文章生成を単語単位でやってみる

ABOUTこの記事をかいた人

NEW POSTこのライターの最新記事

LightWeight_GANで、新しいラーメンを創作してみる

AIで画像から鉛筆画を描く

CPUだけで顔・手・ポーズのリアルタイム検出を行う

ディープラーニング で、1枚の顔画像から詳細3Dモデルを作って自由に動かす

最近の投稿

最近のコメント

アーカイブ

カテゴリー

メタ情報

コメントを残すコメントをキャンセル

ディープラーニングで、1枚の顔画像から詳細3Dモデルを作って自由に動かす