Colabで動く！！Faster R-CNNのPyTorch実装

# はじめに
実は1回目のqiita投稿でFaster-rcnnの実装は出したんですが環境やpathの類が扱いずらいものになってしまったのでcolabで誰でも使えるようにしよう！と思って作りました。
とりあえず物体検出をやってみたい！という方に読んでいただけると幸いです。
ちなみに1回目の記事は[こちら](https://qiita.com/ImR0305/items/22c199f85c44890e0ff8)

## 諸注意
※本記事は**PSCAL VOCフォーマットのデータセット向け**です。
私はBDD100KというデータセットをPascalVOCフォーマットに変換して学習を行ったためclassラベルがBDD100Kのものとなっています。
（PSCAL VOCフォーマットであればpathを変えれば独自データの学習も可能です。）

## コード&簡易データセット
すべてのコードは[github](https://github.com/Ryunosuke-Ikeda/colab_frcnn)に上げます。
データセットの展開後の想定ディレクトリ構造

colab_frcnn-main/
┣Faster-R-CNN.ipynb
┣ smallbdd/
　　　┣test/
　　　┣xml/
　　　┣img/


## localで動かす際の環境構築
コードをクローンしたうえで以下のコマンドを入力してください．

```
pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio==0.8.0 -f https://download.pytorch.org/whl/torch_stable.html
pip install jupyterlab
pip install -r requirements.txt
```
その他バージョン
```
python 3.6
CUDA Toolkit 11.1 Update 1
cuDNN v8.1.1 (Feburary 26th, 2021), for CUDA 11.0,11.1 and 11.2
```



## その１前準備
### ドライブのマウント
仮にローカルで動かす場合はこの操作は必要ありません。

```python

from google.colab import drive
drive.mount('/content/drive')
```

実行すると以下の画面になるので、URL飛んだ先のコードを記入してください。
![スクリーンショット (205)_LI.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/556547/d62e26a4-81e1-6126-1096-4da74588fc7f.jpeg)

### ハイパーパラメータ、pathの指定
仮にオリジナルデータを使う際はここのコード(path)とclassを変えるだけで動くはず！

マイドライブまでのパスである**'/content/drive/My Drive/'**以降にクローンしたcolab_frcnn-mainのフォルダーまでのパスを入力してください。

```python

#pathの指定(colab_frcnn-main直下まで)
path='/content/drive/My Drive/************'
bdd_xml=path+"/colab_frcnn-main/smallbdd/xml"
bdd_img=path+"/colab_frcnn-main/smallbdd/img"
test_path=path+"/colab_frcnn-main/smallbdd/test"

#datasetのクラス指定
dataset_class=['person', 'traffic light', 'train', 'traffic sign', 'rider', 'car', 'bike', 'motor', 'truck', 'bus']
#表示したいラベルの色の指定
#注意！！一番最初は背景クラスを示すので(0,0,0)にする(それ以外は自由)
colors = ((0,0,0),(255,0,0),(0,255,0),(0,0,255),(100,100,100),(50,50,50),(255,255,0),(255,0,255),(0,255,255),(100,100,0),(0,100,100))

#ハイパーパラメータの指定
epochs=10
batch_size=1
scale=720#画像のスケール設定(縦の大きさを入力)

```

### xmlからデータの読み込み
```python
import numpy as np
import pandas as pd
 
from PIL import Image
from glob import glob
import xml.etree.ElementTree as ET 
import cv2
 
import torch
import torchvision
from torchvision import transforms
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torch.utils.data import TensorDataset
import os


class xml2list(object):
    
    def __init__(self, classes):
        self.classes = classes
        
    def __call__(self, xml_path):
        
        ret = []
        xml = ET.parse(xml_path).getroot()
        
        boxes = []
        labels = []
        zz=0
        
        for zz,obj in enumerate(xml.iter('object')):
            
            label = obj.find('name').text
          
            ##指定クラスのみ

            if label in self.classes :
                bndbox = obj.find('bndbox')
                xmin = int(bndbox.find('xmin').text)
                ymin = int(bndbox.find('ymin').text)
                xmax = int(bndbox.find('xmax').text)
                ymax = int(bndbox.find('ymax').text)
                boxes.append([xmin, ymin, xmax, ymax])
                labels.append(self.classes.index(label))
            else:
                continue
            
        num_objs = zz +1

        anno = {'bboxes':boxes, 'labels':labels}

        return anno,num_objs
```
### torchのDatasetを作る
```python
class MyDataset(torch.utils.data.Dataset):
    
        def __init__(self,image_dir,xml_paths,scale,classes):
            
            super().__init__()
            self.image_dir = image_dir
            self.xml_paths = xml_paths
            self.image_ids = sorted(glob('{}/*'.format(xml_paths)))
            self.scale=scale
            self.classes=classes
            
        def __getitem__(self, index):
    
            transform = transforms.Compose([
                                            transforms.ToTensor()
            ])
    
            # 入力画像の読み込み
            image_id=self.image_ids[index].split("/")[-1].split(".")[0]
            image = Image.open(f"{self.image_dir}/{image_id}.jpg")
            
            
            #画像のスケール変換
            t_scale_tate=self.scale ##目標のスケール(縦)
            #縮小比を計算
            ratio=t_scale_tate/image.size[1]
            ##目標横スケールを計算
            t_scale_yoko=image.size[0]*ratio
            t_scale_yoko=int(t_scale_yoko)
            
            #print('縮小前:',image.size)
            #print('縮小率:',ratio)
            #リサイズ
            image = image.resize((t_scale_yoko,t_scale_tate))
            #print('縮小後:',image.size)
  
            image = transform(image)
        
            transform_anno = xml2list(self.classes)
            path_xml=f'{self.xml_paths}/{image_id}.xml'
            

            annotations,obje_num= transform_anno(path_xml)

            boxes = torch.as_tensor(annotations['bboxes'], dtype=torch.int64)
            labels = torch.as_tensor(annotations['labels'], dtype=torch.int64)

          
            #bboxの縮小
            #print('縮小前:',boxes)
            boxes=boxes*ratio
            #print('縮小後:',boxes)
        
            area = (boxes[:, 3]-boxes[:, 1]) * (boxes[:, 2]-boxes[:, 0])
            area = torch.as_tensor(area, dtype=torch.float32)

            iscrowd = torch.zeros((obje_num,), dtype=torch.int64)

            target = {}
            target["boxes"] = boxes
            target["labels"] = labels+1
            target["image_id"] = torch.tensor([index])
            target["area"] = area
            target["iscrowd"] = iscrowd
            return image, target,image_id
        
        def __len__(self):

            return len(self.image_ids)
```
### Dataloaderの作成
```python

def dataloader (data,dataset_class,batch_size,scale=720):
    xml_paths=data[0]
    image_dir1=data[1]
    dataset = MyDataset(image_dir1,xml_paths,scale,dataset_class)

    #データのロード
    torch.manual_seed(2020)
    def collate_fn(batch):
        return tuple(zip(*batch))

    train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True,collate_fn=collate_fn)
    

    return train_dataloader
```
## その2modelの定義
resnet50で学習済みのものを用います。
もしほかのモデルを試したい場合は[1回目の記事](https://qiita.com/ImR0305/items/22c199f85c44890e0ff8#%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E5%AE%9A%E7%BE%A9)のmodel構築部分のmodel2を参考にしてください。


```python

def model ():
    #モデルの定義
    
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
    num_classes=len(dataset_class)+1
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model

```

## その3学習
学習後に学習済みモデルを/colab_frcnn直下にmodel.ptとして保存します。


```python
data_ALL=[bdd_xml,bdd_img]
train_dataloader=dataloader(data_ALL,dataset_class,batch_size,scale)

model=model()
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
num_epochs = epochs



device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 

model.cuda()

model.train()#学習モードに移行

loss_list=[]
for epoch in range(num_epochs):
    loss_epo=[]

    
    for i, batch in enumerate(train_dataloader):
        
 
        images, targets, image_ids = batch#####　batchはそのミニバッジのimage、tagets,image_idsが入ってる
        
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        
        
        ##学習モードでは画像とターゲット（ground-truth）を入力する
        ##返り値はdict[tensor]でlossが入ってる。（RPNとRCNN両方のloss）
        loss_dict= model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        loss_value = losses.item()
        
        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

        #lossの保存
        loss_epo.append(loss_value)
 
        if (i+1) % 10== 0:
          print(f"epoch #{epoch+1} Iteration #{i+1} loss: {loss_value}") 
          
        
    #Epochごとのlossの保存
    loss_list.append(np.mean(loss_epo))
    torch.save(model, path+'/colab_frcnn-main/model.pt')
```

## その4テスト
保存済みの学習済みモデルを使用したい場合は一番下のコメントアウトを外して、一番最初のその1前処理の部分とこれ以降のセルを実行することで、学習部分を実行しなくてもテストが可能です。

```python
import cv2
import glob
import matplotlib.pyplot as plt

data_class=dataset_class
data_class.insert(0, "__background__")
classes = tuple(data_class)

#学習済みモデルで推論する場合
#model=torch.load(path+'/colab_frcnn-main/model.pt')
```
### 結果の表示
```python
model.to(device)

model.eval()
for imgfile in sorted(glob.glob(test_path+'/*')):
    
    img = cv2.imread(imgfile)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    image_tensor = torchvision.transforms.functional.to_tensor(img)

    with torch.no_grad():
        prediction = model([image_tensor.to(device)])
    
    for i,box in enumerate(prediction[0]['boxes']):
        score = prediction[0]['scores'][i].cpu().numpy()
        if score > 0.5:
            score = round(float(score),2)
            cat = prediction[0]['labels'][i].cpu().numpy()
            txt = '{} {}'.format(classes[int(cat)], str(score))
            font = cv2.FONT_HERSHEY_SIMPLEX
            cat_size = cv2.getTextSize(txt, font, 0.5, 2)[0]
            c = colors[int(cat)]
            box=box.cpu().numpy().astype('int')
            cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), c , 2)
            cv2.rectangle(img,(box[0], box[1] - cat_size[1] - 2),(box[0] + cat_size[0], box[1] - 2), c, -1)
            cv2.putText(img, txt, (box[0], box[1] - 2), font, 0.5, (0, 0, 0), thickness=1, lineType=cv2.LINE_AA)


    plt.figure(figsize=(15,10))
    plt.imshow(img)
plt.show()
```
##結果
こんな画像が出力されるはず
![ダウンロード.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/556547/22a28c15-55df-97db-c0ed-5115edd41b40.png)

##終わりに
いかがだったでしょうかこれを使えばcolabで簡単に実行ができた(はず)です。
またpathを変えればオリジナルデータセット(PASCAL VOCフォーマット)での学習も簡単にできます。
Faster-R-CNNの実装に苦しんでいる、実際に使ってみたかっ方の参考になれば幸いです。