今回は、血液の顕微鏡画像から細胞を検出する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ネットワークが出来ます。早速、コードを書いてみましょう。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 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 に変更して下さい。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 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 に変更します。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
# 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）を使って、テスト画像の推論を実行します。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
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ネットワークなので、結構物体検出ができるものですね。
では、また。
コメントを残す