Pose grammar training

training.py

import numpy as np
from pycocotools.coco import COCO
import os
import math
import keras
from keras.models import Model
from keras.layers import Dense, Input, Flatten, Reshape
from keras.utils import plot_model
from random import shuffle

width = 18
height = 28
thres = 0.05
coco_train = COCO(os.path.join('/media/muhammed/Other/RESEARCH/datasets/MSCOCO/coco/annotations',
                               'person_keypoints_train2017.json'))
coco_val = COCO(os.path.join('/media/muhammed/Other/RESEARCH/datasets/MSCOCO/coco/annotations',
                             'person_keypoints_val2017.json'))
batch_size = 128

def get_data(ann_data, coco):
    weights = np.zeros((width, height, 17))
    output = np.zeros((width, height, 17))

    bbox = ann_data['bbox']
    x0 = int(bbox[0])
    y0 = int(bbox[1])
    x = float(bbox[2])
    y = float(bbox[3])

    xscale = float(width) / math.ceil(x)
    yscale = float(height) / math.ceil(y)

    kpx = ann_data['keypoints'][0::3]
    kpy = ann_data['keypoints'][1::3]
    kpv = ann_data['keypoints'][2::3]

    for xy in range(17):
        if kpv[xy] > 0:
            a = int(round(kpx[xy] - x0) * xscale)
            b = int(round(kpy[xy] - y0) * yscale)

            if a >= width and b >= height:
                weights[width - 1][height - 1][xy] = 1
            elif a >= width:
                weights[width - 1][b][xy] = 1
            elif b >= height:
                weights[a][height - 1][xy] = 1
            else:
                output[a, b, xy] = 1

    img_id = ann_data['image_id']
    img_data = coco.loadImgs(img_id)[0]
    ann_data = coco.loadAnns(coco.getAnnIds(img_data['id']))

    for a in range(len(ann_data)):
        kpx = ann_data[a]['keypoints'][0::3]
        kpy = ann_data[a]['keypoints'][1::3]
        kpv = ann_data[a]['keypoints'][2::3]

        for x in range(17):
            if kpv[x] > 0:
                if (kpx[x] > bbox[0] - bbox[2] * thres and kpx[x] < bbox[0] + bbox[2] * (1 + thres)):
                    if (kpy[x] > bbox[1] - bbox[3] * thres and kpy[x] < bbox[1] + bbox[3] * (1 + thres)):
                        a = int(round(kpx[x] - x0) * xscale)
                        b = int(round(kpy[x] - y0) * yscale)
                        if a >= width and b >= height:
                            weights[width - 1][height - 1][x] = 1
                        elif a >= width:
                            weights[width - 1][b][x] = 1
                        elif b >= height:
                            weights[a][height - 1][x] = 1
                        elif a < width and b < height:
                            weights[a][b][x] = 1
    return weights, output

def train_bbox_generator():
    ann_ids = coco_train.getAnnIds()
    while 1:
        shuffle(ann_ids)
        for i in range(len(ann_ids) // batch_size):
            X = np.zeros((batch_size, width, height, 17))
            Y = np.zeros((batch_size, width, height, 17))
            for j in range(batch_size):
                ann_data = coco_train.loadAnns(ann_ids[i + j])[0]
                x, y = get_data(ann_data, coco_train)
                X[j, :, :, :] = x
                Y[j, :, :, :] = y
            yield X, Y

def val_bbox_generator():
    ann_ids = coco_val.getAnnIds()
    while 1:
        shuffle(ann_ids)
        for i in range(len(ann_ids) // batch_size):
            X = np.zeros((batch_size, width, height, 17))
            Y = np.zeros((batch_size, width, height, 17))
            for j in range(batch_size):
                ann_data = coco_val.loadAnns(ann_ids[i + j])[0]
                x, y = get_data(ann_data, coco_val)
                X[j, :, :, :] = x
                Y[j, :, :, :] = y
            yield X, Y

input = Input(shape=(width,height,17))
x = Flatten()(input)
x = Dense(width*height*17, activation='relu')(x)
x = Dense(50, activation='relu')(x)
x = Dense(width*height*17, activation='softmax')(x)
x = Reshape((width,height,17))(x)

model = Model(inputs=input, outputs=x)

print(model.summary())

plot_model(model, to_file='test.png', show_shapes=True)

adam_optimizer = keras.optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)

model.compile(loss='binary_crossentropy', optimizer=adam_optimizer)

checkpoint = keras.callbacks.ModelCheckpoint('weights.{epoch:02d}-{val_loss:.2f}.h5',verbose=1)
csv_log    = keras.callbacks.CSVLogger('training_log.csv', separator=',', append=False)

model.fit_generator(generator = train_bbox_generator(),
                    steps_per_epoch = len(coco_train.getAnnIds()) // batch_size,
                    validation_data= val_bbox_generator(),
                    validation_steps= len(coco_val.getAnnIds()) // batch_size,
                    epochs = 50,
                    callbacks=[checkpoint, csv_log],
                    verbose=1)