PyTorchでPascal VOCデータセット を使ってみた
はじめに
Object DetectionのデータセットPascal VOCデータセットをPyTorchで使ってみた。
Pascal VOCとは
Pattern Analysis, Statistical Modelling and Computational Learning(PASCAL) Visual Object Classes Challenge(VOC)。 2005年～2012年に開催されていた物体のクラス認識などのコンペティション。 物体検出やセグメンテーションなどのタスクがある。
クラス数は2007年からは20クラス。
- Person: person
- Animal: bird, cat, cow, dog, horse, sheep
- Vehicle: aeroplane, bicycle, boat, bus, car, motorbike, train
- Indoor: bottle, chair, dining table, potted plant, sofa, tv/monitor
その中で今回は、物体検出のデータセットを使ってみた。
VOCDetectionクラス
PyTorchでPascal VOCのデータセットを使うのに、 Torchvisionの VOCDetection というクラスがある。 Pascal VOCのデータを簡単に読み込めるようなクラスが用意されている。 これを使ってみた。
クラスは下のような引数をとる。
torchvision.datasets.VOCDetection(root, year='2012', image_set='train', download=False, transform=None, target_transform=None, transforms=None)
torch.utils.data.Datasetのサブクラスで、
データを得る__getitem__と、
データセットの大きさを得る __len__を持っている。
引数は下のようなもの。
| 引数名 | 説明 | 値 |
|---|---|---|
| root | データセットのルートフォルダー | (e.g.) ./VOCDetection |
| year | どの年のデータセットか | 2007, 2008, 2009, 2010, 2011, 2012 |
| image_set | データセットの区分 | train, trainval, val |
| download | データセットをダウンロードするかどうか | True, False |
| transform | 画像に対して変形を行う際のtransform関数 | None(何もしない), torchvision.transforms.ToTensor()など |
| target_transform | targetに対して変形を行う際のtransform関数 | None(何もしない) など |
| transforms | transformとtarget_transformを合わせたようなやつ。？なんで必要なのこれ？ | None(何もしない)など |
year=2007の時のみ、image_set="test"も使える。
indexでアクセスできる。 返り値は、(image,target)で、imageはPILの画像、targetはxml treeの辞書となっている。 targetには画像サイズやbounding boxの位置などの情報が入っている。
試してみた
バージョン
PyTorch: 1.4.0
Torchvision : 0.5.0
Google Colabで実行
まずは、データセットクラスを宣言。2012年のデータセットのtrainを使う。
download=Trueなので、実行するとデータがrootで指定した./VOCDetection/2012にダウンロードされる。
import torchvision voc_dataset=torchvision.datasets.VOCDetection(root="./VOCDetection/2012",year="2012",image_set="train",download=True)
データセットから一番最初のデータを取得する。 返り値は、画像とアノテーション情報である。
image,target=voc_dataset[0]
データセットの画像を表示してみる。
from IPython.display import display display(image)
次に、アノテーション情報targetについて見てみる。 targetを表示してみると、下のような辞書の形でアノテーション情報が入っている。
{'annotation': {'folder': 'VOC2012', 'filename': '2008_000008.jpg', 'source': {'database': 'The VOC2008 Database', 'annotation': 'PASCAL VOC2008', 'image': 'flickr'}, 'size': {'width': '500', 'height': '442', 'depth': '3'}, 'segmented': '0', 'object': [{'name': 'horse', 'pose': 'Left', 'truncated': '0', 'occluded': '1', 'bndbox': {'xmin': '53', 'ymin': '87', 'xmax': '471', 'ymax': '420'}, 'difficult': '0'}, {'name': 'person', 'pose': 'Unspecified', 'truncated': '1', 'occluded': '0', 'bndbox': {'xmin': '158', 'ymin': '44', 'xmax': '289', 'ymax': '167'}, 'difficult': '0'}]}}
画像のサイズ、物体の名前と位置などを使うので取り出してみる。 bounding boxの情報は、画像の一番左上のピクセルを(1,1)とした座標であり、左上の頂点と右下の頂点の情報で保持されている。
annotation=target["annotation"] #画像サイズ width, height, depth = annotation["size"].values() #物体の情報 objects=annotation["object"] for obj in objects: name=obj["name"] bndbox=obj["bndbox"] #bounding boxの左上と右下の座標 x_min,y_min,x_max,y_max=[int(val) for val in bndbox.values()]
画像とbounding boxが得られたので、画像上にbounding boxを表示してみる。 bounding boxの適用した画像を作る関数をつくる。もしかしたら、すでに便利な関数があるかもしれないが、 少し調べた範囲ではわからなかった。
from PIL import Image,ImageDraw def make_bbox_img(pil_image,x0y0,x1y1,text,width=1,color="red"): bbox_img=pil_image.copy() d=ImageDraw.Draw(bbox_img) d.rectangle((x0y0,x1y1),outline=color,width=width) #テキストのライン d.text((x0y0[0]+1,x0y0[1]+1),text,fill="white") d.text((x0y0[0]-1,x0y0[1]-1),text,fill="white") #テキスト d.text(x0y0,text,fill="black") return bbox_img name_to_color={ "person":"orange", "bird":"darkcyan", "cat":"lightyellow", "cow":"coral", "dog":"dimgray", "horse":"darkgreen", "sheep":"wheat", "aeroplane":"red", "bicycle":"royalblue", "boat":"blue", "bus":"yellow", "car":"hotpink", "motorbike":"lawngreen", "train":"cyan", "bottle":"olive", "chair":"indigo", "diningtable":"maroon", "pottedplant":"mediumslateblue", "sofa":"purple", "tvmonitor":"lime", } def make_voc_bbox_img(pil_image,objects,width=1): bbox_img=pil_image.copy() d=ImageDraw.Draw(bbox_img) #list形式でないものがあるのでlist形式へ直す if not isinstance(objects,list): objects=[objects] for obj in objects: name=obj["name"] bndbox=obj["bndbox"] x_min,y_min,x_max,y_max=[int(val) for val in bndbox.values()] color=name_to_color[name] #x:[1,image_width],y:[1,image_height] bbox_img=make_bbox_img(bbox_img,(x_min-1,y_min-1),(x_max-1,y_max-1),text=name,color=color,width=width) return bbox_img
先ほどの画像に適用してみる。
bbox_img=make_voc_bbox_img(image,objects) display(bbox_img)
うまくbounding boxを表示出来ている。 色が見にくいなどはあるが、表示できるようになったので、まあよしとする。 bounding box表示方法は、簡単な方法がないか探していきたい。
数枚の画像で試してみた。
つまづいた点
つまづいた点としては下のものがある。
- annotationのbounding boxの座標の値がstr型
- annotationのobjectの値が、全てlist形式と思っていたが、
objectが複数あるものはlist形式で、単一のものはlist形式でない点。
まとめ
PyTorchで物体検出を試してみたくて、Pascal VOCデータセットについて調べて動かしてみた。 データ形式などがわかった。 すでにクラスが用意されているので、自分でデータセットの読み込みなどを書く必要がなくて便利。 bounding box表示の関数がないか今後調べていきたい。
