<jupyter_start><jupyter_text>Pickle files format:the processed_bodyinfo is a numpy array that gets dumped into a pickle is in the following format:bodyinfo[f][b][j][u]Where:     f is the frame number             (depends on length of the video)    b is the body number              (there might not always be just 1 body/skeleton in a frame)    j is the joint number for a body  (25)    u is the basis number for a joint (3)<jupyter_code>pickle_files = glob.glob(skeletons_dir+'*.pickle')
with open(pickle_files[0], 'rb') as fp:
    bodyinfo = pickle.load(fp)

bodyinfo.shape
# Vertices to keep:
idxKeep = [3,4,5,6,7,8,9,10,11,21,22,23,24]
h1 = util.pickle_to_action_array(pickle_files[0], idxKeep)
h1[0][1].shape
# Velocity
h1[0][1][0:20]
# Accelaration
h1[0][2][:,4,:]
# x-axis
plt.hist(h1[0][2][:,4,0])
# y-axis -hand
plt.hist(h1[0][2][:,4,1])
# y-axis -head
plt.hist(h1[0][2][:,0,1])
# z-axis
plt.hist(h1[0][2][:,4,2])
a_all = np.asarray([x[0][2] for x in data]).flatten()
#a values:
plt.hist(a_all, bins =50)
data = util.pickle_files_to_action_series(pickle_files, idxKeep)<jupyter_output>./data_raw/S002C003P010R002A010.pickle
./data_raw/S003C003P002R002A010.pickle
./data_raw/S004C002P008R001A010.pickle
./data_raw/S005C002P016R001A010.pickle
./data_raw/S006C003P007R002A010.pickle
./data_raw/S007C002P026R001A010.pickle
./data_raw/S007C003P027R002A010.pickle
./data_raw/S008C002P001R001A010.pickle
./data_raw/S011C002P027R001A010.pickle<jupyter_text>Choose hI to be (-0.6, -0.4, -0.2, 0, 0.2)<jupyter_code>with open('./data_raw/S004C002P008R001A010.pickle', 'rb') as fp:
    bodyinfo_check = pickle.load(fp)

bodyinfo_check.shape
# Those seem to be broken files
len(data)
data[0][0].shape<jupyter_output><empty_output><jupyter_text>Visualize normalized state<jupyter_code>state = data[0][0]
frame = 1
xs_1 = state[frame,[0,1,2,3,4,9,10] ,0]
xs_2 = state[frame,[0,5,6,7,8,11,12],0]
ys_1 = state[frame,[0,1,2,3,4,9,10],1] 
ys_2 = state[frame,[0,5,6,7,8,11,12],1]
zs_1 = state[frame,[0,1,2,3,4,9,10],2] 
zs_2 = state[frame,[0,5,6,7,8,11,12],2]
x_max, x_min = np.max(state[:,:,0]), np.min(state[:,:,0])
y_max, y_min = np.max(state[:,:,1]), np.min(state[:,:,1])
z_max, z_min = np.max(state[:,:,2]), np.min(state[:,:,2])
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_xlim3d(x_min-1e-3, x_max+1e-3)
ax.set_ylim3d(y_min-1e-3, y_max+1e-3)
ax.set_zlim3d(z_min-1e-3, z_max+1e-3)
ax.plot(xs_1, ys_1, zs_1, marker = '^')
ax.plot(xs_2, ys_2, zs_2, marker = '^')
fig.savefig('test.png')
def getImage(data, key, state =None):
    if state is None:
        state = data[key][0]
    x_max, x_min = np.max(state[:,:,0]), np.min(state[:,:,0])
    y_max, y_min = np.max(state[:,:,1]), np.min(state[:,:,1])
    z_max, z_min = np.max(state[:,:,2]), np.min(state[:,:,2])
    path = './plot_' + str(key)
    images = []
    if not os.path.isdir(path):
        os.mkdir(path)
    for frame in range(len(state)):
        xs_1 = state[frame,[0,1,2,3,4,9,10] ,0]
        xs_2 = state[frame,[0,5,6,7,8,11,12],0]
        ys_1 = state[frame,[0,1,2,3,4,9,10],1] 
        ys_2 = state[frame,[0,5,6,7,8,11,12],1]
        zs_1 = state[frame,[0,1,2,3,4,9,10],2] 
        zs_2 = state[frame,[0,5,6,7,8,11,12],2]

        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.set_xlim3d(x_min-1e-3, x_max+1e-3)
        ax.set_ylim3d(y_min-1e-3, y_max+1e-3)
        ax.set_zlim3d(z_min-1e-3, z_max+1e-3)
        ax.plot(xs_1, ys_1, zs_1, marker = '^')
        ax.plot(xs_2, ys_2, zs_2, marker = '^')
        fig.savefig(path + '/plot_' + str(frame) + '.png')
        plt.close(fig)
        images.append(imageio.imread(path + '/plot_' + str(frame) + '.png'))
    imageio.mimsave(path + '/movie.gif', images)
getImage(data, 0)
len(data[100][0])
getImage(data, 971)
data[100][-1]
pickle.dump(data, open('data.p', 'wb'))<jupyter_output><empty_output><jupyter_text>Build model<jupyter_code>import torch
import torch.nn as nn
from torch.autograd import Variable
import pickle
import torch.nn.functional as F
from torch.utils.data import Dataset
from torch import optim
import random
gamma = 0.95
def getReward(seq, gamma):
    len_seq = len(seq)+1
    idx = np.arange(len_seq)
    reward = gamma**(len_seq - idx-1)
    reward = np.diff(reward)
    return reward
a = [1,2,3,4,5]
getReward(a, gamma)
import pdb
class skDataset(Dataset):
    def __init__(self, data, gamma, length):
        self.data = data
        self.gamma = gamma
        self.length = length

    def __len__(self):
        return len(self.data)

    def __getitem__(self, key):
        item = self.data[key]
        v_seq = item[0][1:-1]
        d_seq = item[1][0:-1]
        a_seq = item[2]
        reward = getReward(a_seq, self.gamma)
        
        pad_len = max(0, self.length - len(v_seq))
        v_seq_pad = np.pad(v_seq, pad_width = ((0, pad_len),(0,0), (0,0)), mode = 'constant', constant_values=0)
        d_seq_pad = np.pad(d_seq, pad_width = ((0, pad_len),(0,0), (0,0)), mode = 'constant', constant_values=0)
        a_seq_pad = np.pad(d_seq, pad_width = ((0, pad_len),(0,0), (0,0)), mode = 'constant', constant_values=0)
        a_seq_pad = np.clip(a_seq_pad, -2,2)
        reward_pad = np.pad(reward, pad_width = (0, pad_len), mode = 'constant', constant_values=0)
        
        sample = {'v_seq': torch.from_numpy(v_seq_pad).float(),
                  'd_seq': torch.from_numpy(d_seq_pad).float(),
                  'a_seq': torch.from_numpy(a_seq_pad).float(),
                  'reward': torch.from_numpy(reward_pad).float()}
        return sample

lsLen = [len(x[0]) for x in data]
print(np.percentile(lsLen, 90))
print(np.percentile(lsLen, 95))
print(np.max(lsLen))
train = skDataset(data[:-10], gamma, 140)
test = skDataset(data[-10:], gamma, 140)
class gru(nn.Module):
    def __init__(self, dimIn, dimLSTM, nLSTM, p_dropOut, lsDim, flg_bn):
        super(gru, self).__init__()
        self.dimLSTM = dimLSTM
        self.nLSTM = nLSTM
        self.p_dropOut = p_dropOut
        self.flg_bn = flg_bn
        self.lsDim = lsDim
        
        self.lstm = nn.GRU(dimIn, dimLSTM, self.nLSTM, batch_first=True, bidirectional=False, dropout=self.p_dropOut)
        self.bns = nn.ModuleList([nn.BatchNorm1d(self.lsDim[i + 1]) for i in range(len(self.lsDim) - 2)])
        self.FCs = nn.ModuleList([nn.Linear(self.lsDim[i], self.lsDim[i + 1]) for i in range(len(self.lsDim) - 1)])

    def fc_layer(self, x, layer, bn=None):
        x = layer(x)
        if self.flg_bn:
            x = F.relu(bn(x))
        else:
            x = F.relu(x)
        x = F.dropout(x, self.p_dropOut)
        return x

    def init_hidden(self, batchSize, nLSTM):
        return Variable(torch.zeros(nLSTM, batchSize, self.dimLSTM))

    def forward(self, v_seq, d_seq):
        batchSize, seq_len = v_seq.shape[0], v_seq.shape[1]
        in_seq = torch.cat((v_seq, d_seq), 2).view(batchSize, seq_len, -1)
        h0 = self.init_hidden(batchSize=batchSize, nLSTM= self.nLSTM)
        lstm_out, hidden = self.lstm(in_seq, h0)

        z = lstm_out
        if len(self.lsDim) > 2:
            for i in range(len(self.lsDim) - 2):
                z = self.fc_layer(z, self.FCs[i], self.bns[i])

        self.y_hat = 2*F.tanh(self.FCs[-1](z))
        return self.y_hat

    def prediction(self, v_1, d_0, max_len):
        batchSize = v_1.shape[0]
        v_seq = torch.zeros(batchSize, max_len+1, 13, 3)
        v_seq[:,0,:,:] = v_1
        d_seq = torch.zeros(batchSize, max_len+1, 13, 3)
        d_seq[:,0,:,:] = d_0
        a_seq = torch.zeros(batchSize, max_len, 13, 3)
        in_seq = torch.cat((v_1, d_0), 2).view(batchSize, 1, -1)
        hidden = self.init_hidden(batchSize = batchSize, nLSTM= self.nLSTM)
        for i in range(max_len):
            lstm_out, hidden = self.lstm(in_seq, hidden)
            z = lstm_out
            if len(self.lsDim) > 2:
                for i in range(len(self.lsDim) - 2):
                    z = self.fc_layer(z, self.FCs[i], self.bns[i])
            y_hat = 2*F.tanh(self.FCs[-1](z))
            a_i = y_hat.view(batchSize, -1, 3)
            a_seq[:,i,:,:] = a_i
            d_seq[:,i+1,:,:] = d_seq[:,i,:,:] + a_seq[:,i,:,:]
            v_seq[:,i+1,:,:] = d_seq[:,i+1,:,:] + v_seq[:,i,:,:]
            in_seq = torch.cat((v_seq[:,(i+1):(i+2),:,:], d_seq[:,(i+1):(i+2),:,:]), 2).view(batchSize,1, -1)
        return v_seq, d_seq, a_seq
batchSize = 32
train_loader = torch.utils.data.DataLoader(train, batch_size=batchSize, shuffle=True, pin_memory=True)
test_loader = torch.utils.data.DataLoader(test, batch_size=batchSize, shuffle=False, pin_memory=True)
model = gru(dimIn = 78, dimLSTM = 128, nLSTM = 1, p_dropOut = 0.5, lsDim = [128,39], flg_bn = False)
lr = 0.001
decay_per_epoch = 10
criterion = torch.nn.MSELoss(reduction='none')
opt = optim.Adam(model.parameters(), lr=lr)
n_batch = 20
tf = 0.7
tf=0.5
for epoch in range(200):
    #--- Training part ------
    train_loss = 0
    model.train()
    lr = max(lr * (0.9 ** (epoch // decay_per_epoch)), 1e-5) 
    for param_group in opt.param_groups:
        param_group['lr'] = lr
    
    use_tf = True if random.random() < tf else False
    j, nRec_train = 0, 0
    while j <= n_batch: # Use this to adjust how many batches as one epoch
        # for batch_idx, sample in enumerate(self.train_loader):
        sample = train_loader.__iter__().__next__()
        v_seq, d_seq, a_seq, reward = sample['v_seq'], sample['d_seq'], sample['a_seq'], sample['reward']
        batchSize_j, seq_len = v_seq.shape[0], v_seq.shape[1]
        v_seq, d_seq, a_seq, reward = Variable(v_seq).float(), Variable(d_seq).float(), Variable(a_seq).float(), \
        Variable(reward).float()
        nRec_j = (reward != 0).sum().float()
        nRec_train += nRec_j

        opt.zero_grad()
        if use_tf:
            a_hat = model(v_seq, d_seq)
            loss = criterion(a_hat, a_seq.view(batchSize_j, seq_len, -1))
            loss = (loss * reward.view(batchSize_j, seq_len,1).repeat(1,1,loss.shape[-1])).sum() / nRec_j*100
        else:
            v_seq_pred, d_seq_pred, a_seq_pred  = model.prediction(v_seq[:,0,:,:], d_seq[:,0,:,:], max_len=140)
            loss = criterion(a_seq_pred, a_seq)
            loss = (loss * reward.view(batchSize_j, seq_len,1,1).repeat(1,1,13,3)).sum() / nRec_j*100
        loss.backward()
        opt.step()
        j += 1

        train_loss += loss.data.item() * nRec_j
        
    #----- Test part ------
    model.eval()
    test_loss, nRec = 0, 0
    for batch_idx, sample in enumerate(test_loader):
        v_seq, d_seq, a_seq, reward = sample['v_seq'], sample['d_seq'], sample['a_seq'], sample['reward']
        batchSize_j, seq_len = v_seq.shape[0], v_seq.shape[1]
        nRec_j = (reward != 0).sum().float()
        nRec += nRec_j
        v_seq, d_seq, a_seq, reward = Variable(v_seq).float(), Variable(d_seq).float(), Variable(a_seq).float(), \
        Variable(reward).float()
        
        v_seq_pred, d_seq_pred, a_seq_pred  = model.prediction(v_seq[:,0,:,:], d_seq[:,0,:,:], max_len=140)
        test_loss_step = criterion(a_seq_pred, a_seq)
        test_loss_step = (test_loss_step * reward.view(batchSize_j, seq_len,1,1).repeat(1,1,13,3)).sum() / nRec_j
        test_loss += test_loss_step.data.item()*nRec_j*100

    test_loss /= nRec
    print('Use TF {}, Train Epoch: {}, Loss: {:.4f}, Test loss: {:.4f}'.format(use_tf, epoch, train_loss/nRec_train, test_loss))

## Generation
i = 9
sample = test[i]
v_seq, d_seq, a_seq,reward = sample['v_seq'], sample['d_seq'], sample['a_seq'],sample['reward']
v_seq.shape
list(reward.numpy()).index(1) # Take last frame as init
a_seq[56:60,0,0]
v1 = v_seq[56,:,:].view(1,13,3)
d0 = d_seq[56,:,:].view(1,13,3)
model.eval()
v_seq_pred, d_seq_pred, a_seq_pred  = model.prediction(v1, d0, max_len=40)
v_seq_pred[0,9,:,:]
a_seq_pred[0,5,:,:]
a_seq[52,:,:]
v_seq[56,:,:]
v_seq_pred.shape
dataPlot = v_seq_pred[0].data.numpy()
getImage(None, key='test_9', state = dataPlot[0:4,:,:])
getImage(None, key = 'test_9_orig', state = v_seq[0:56,:,:].data.numpy())<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1789229878401853, 'fraction_numerical': 0.05525684506576226, 'mean_word_length': 3.2979025950942056, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 18, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '363002', 'license': 'MIT', 'n_tokens_mistral': 5399, 'n_tokens_neox': 5091, 'n_words': 896, 'path': 'HW2.ipynb', 'repo_name': 'tomxi/csci3033', 'size': '12089'}