Regression with BIWI head pose dataset¶

This is a more advanced example to show how to create custom datasets and do regression with images. Our task is to find the center of the head in each image. The data comes from the BIWI head pose dataset, thanks to Gabriele Fanelli et al. We have converted the images to jpeg format, so you should download the converted dataset from this link.

In [ ]:

%matplotlib inline

In [ ]:

from fastai2.vision.all import *
from nbdev.showdoc import *

Getting and converting the data¶

In [ ]:

path = untar_data(URLs.BIWI_HEAD_POSE)

In [ ]:

cal = np.genfromtxt(path/'01'/'rgb.cal', skip_footer=6); cal

Out[ ]:

array([[517.679,   0.   , 320.   ],
       [  0.   , 517.679, 240.5  ],
       [  0.   ,   0.   ,   1.   ]])

In [ ]:

fname = '09/frame_00667_rgb.jpg'

In [ ]:

def img2txt_name(f): return path/f'{str(f)[:-7]}pose.txt'

In [ ]:

img = PILImage.create(path/fname)
img.show();

No description has been provided for this image

In [ ]:

ctr = np.genfromtxt(img2txt_name(fname), skip_header=3); ctr

Out[ ]:

array([187.332 ,  40.3892, 893.135 ])

In [ ]:

def convert_biwi(coords):
    c1 = coords[0] * cal[0][0]/coords[2] + cal[0][2]
    c2 = coords[1] * cal[1][1]/coords[2] + cal[1][2]
    return tensor([c1,c2])

def get_ctr(f):
    ctr = np.genfromtxt(img2txt_name(f), skip_header=3)
    return convert_biwi(ctr)

def get_ip(img,pts): return TensorPoint.create(pts, img_size=img.size)

In [ ]:

get_ctr(fname)

Out[ ]:

tensor([428.5814, 263.9104])

In [ ]:

ctr = get_ctr(fname)
ax = img.show(figsize=(6, 6))
get_ip(img, ctr).show(ctx=ax);

Creating a dataset¶

In [ ]:

dblock = DataBlock(blocks=(ImageBlock, PointBlock),
                   get_items=get_image_files,
                   splitter=FuncSplitter(lambda o: o.parent.name=='13'),
                   get_y=get_ctr,
                   batch_tfms=[*aug_transforms(size=(120,160)), Normalize.from_stats(*imagenet_stats)])

In [ ]:

dls = dblock.dataloaders(path, path=path, bs=64)

In [ ]:

dls.show_batch(max_n=9, figsize=(9,6))

Train model¶

In [ ]:

#TODO: look in after_item for c
dls.c = dls.train.after_item.c

In [ ]:

learn = cnn_learner(dls, resnet34)

In [ ]:

learn.lr_find()

In [ ]:

lr = 2e-2

In [ ]:

learn.fit_one_cycle(5, slice(lr))

epoch	train_loss	valid_loss	time
0	0.140113	0.143032	00:28
1	0.039437	0.002323	00:25
2	0.016588	0.001802	00:25
3	0.009252	0.000881	00:26
4	0.006983	0.001185	00:26

In [ ]:

learn.save('stage-1')

In [ ]:

learn.load('stage-1');

In [ ]:

learn.show_results(max_n=6)

Data augmentation¶

In [ ]:

def repeat_one_file(path): 
    items = get_image_files(path)
    return [items[0]] * 500

In [ ]:

dblock = DataBlock(blocks=(ImageBlock, PointBlock),
                   get_items=repeat_one_file,
                   splitter=RandomSplitter(),
                   get_y=get_ctr)

In [ ]:

tfms = aug_transforms(max_rotate=20, max_zoom=1.5, max_lighting=0.5, max_warp=0.4, p_affine=1., p_lighting=1., size=(120,160))

In [ ]:

dls = dblock.dataloaders(path, path=path, bs=64, batch_tfms=[*tfms, Normalize.from_stats(*imagenet_stats)])

In [ ]:

dls.show_batch(max_n=9, figsize=(8,6))

In [ ]: