def train_cnn():
modeltype = 'cnn'
train_batch_size = 2
sequence_length = 1
directory_type = 'DPM'
pixel = Pixel('pixel.pkl')
train_dataloader, train_data_size, val_dataloader, val_data_size = \
get_train_data(train_batch_size, sequence_length, directory_type, pixel)
cnnmodel = CNN()
if torch.cuda.is_available(): cnnmodel = cnnmodel.cuda()
for i, batch in enumerate(train_dataloader):
frame1 = Variable(torch.stack(batch[0])).squeeze(1)
frame2 = Variable(torch.stack(batch[1]))
labels = Variable(torch.LongTensor(batch[2]))
print(frame1.size())
print(frame2.size())
print(labels)
print(cnnmodel(frame1).size())
break
import torch.nn.functional as functional
import torch.optim as optim
from datautils import Pixel
from pretrain_utils import dataset, get_train_data, ResNetPretraining
use_pickle = False
num_epochs = 300
learning_rate = 0.1
train_batch_size = 1 # tried 64 originally, but memory errors if batch_size >= 4...
sequence_length = 1
directory_type = 'DPM'
modeltype = 'resnet_pretrained'
pixel = Pixel('pixel.pkl')
if use_pickle: pixel.load()
print('Building the train loader...')
train_dataloader, train_data_size, val_dataloader, val_data_size = \
get_train_data(train_batch_size, sequence_length, directory_type, pixel)
save_dir = 'models/{}/'.format(modeltype)
valfile = open('val.txt', 'a')
dict_args = {'input_dim' : 500, 'feature_dim' : 100}
resnet = ResNetPretraining(dict_args)
criterion = nn.NLLLoss()
optimizer = optim.Adadelta(resnet.parameters(), lr=learning_rate, rho=0.95, eps=1e-06, weight_decay=0)
def train_motion(use_pickle=True):
modeltype = 'motion'
train_batch_size = 64
sequence_length = 6
directory_type = 'DPM'
pixel = Pixel('pixel.pkl')
if use_pickle: pixel.load()
pretrained = PreTrainedResnet({'intermediate_layers': ['fc']})
if torch.cuda.is_available():
pretrained = pretrained.cuda()
train_dataloader, train_data_size, val_dataloader, val_data_size = \
get_train_data(train_batch_size, sequence_length, directory_type, pixel, pretrained)
save_dir = 'models/{}/'.format(modeltype)
valfile = open('val.txt', 'a')
dict_args = {
'input_dim': 4,
'rnn_hdim': 200,
'rnn_type': 'LSTM',
'feature_dim': 100,
}
motmodel = MotionModel(dict_args)
num_epochs = 300
learning_rate = 1.0
criterion = nn.NLLLoss()
optimizer = optim.Adadelta(motmodel.parameters(),
lr=learning_rate,
rho=0.95,
eps=1e-06,
weight_decay=0)
if torch.cuda.is_available():
motmodel = motmodel.cuda()
criterion = criterion.cuda()
for epoch in range(num_epochs):
for i, batch in enumerate(train_dataloader):
trackcoords = Variable(torch.stack(batch[3]))
detectioncoord = Variable(torch.stack(batch[4]))
labels = Variable(torch.LongTensor(batch[2]))
if torch.cuda.is_available():
trackcoords = trackcoords.cuda()
detectioncoord = detectioncoord.cuda()
labels = labels.cuda()
motmodel = motmodel.train()
output, _ = motmodel(trackcoords, detectioncoord)
output = functional.log_softmax(output, dim=-1)
loss = criterion(output, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
nn.utils.clip_grad_norm(motmodel.parameters(), 5.0)
if ((i + 1) % 10 == 0):
print('Epoch: [{0}/{1}], Step: [{2}/{3}], Loss: {4}'.format( \
epoch+1, num_epochs, i+1, train_data_size//train_batch_size, loss.data[0]))
#Saving the model
if (epoch % 1 == 0):
if not os.path.isdir(
os.path.join(save_dir, "epoch{}".format(epoch))):
os.makedirs(os.path.join(save_dir, "epoch{}".format(epoch)))
filename = modeltype + '.pth'
file = open(
os.path.join(save_dir, "epoch{}".format(epoch), filename),
'wb')
torch.save(
{
'state_dict': motmodel.state_dict(),
'dict_args': dict_args
}, file)
print('Saving the model to {}'.format(save_dir +
"epoch{}".format(epoch)))
file.close()
#Validation accuracy
if (epoch % 1 == 0):
num_correct = 0.0
num_total = 0.0
for j, batch in enumerate(val_dataloader):
trackcoords = Variable(torch.stack(batch[3]))
detectioncoord = Variable(torch.stack(batch[4]))
labels = Variable(torch.LongTensor(batch[2]))
if torch.cuda.is_available():
trackcoords = trackcoords.cuda()
detectioncoord = detectioncoord.cuda()
labels = labels.cuda()
motmodel = motmodel.eval()
output, _ = motmodel(trackcoords, detectioncoord)
predictions = output.max(dim=-1)[1]
num_correct += (predictions == labels).sum().data[0]
num_total += len(labels)
accuracy = float(num_correct / num_total)
print('Epoch {} Accuracy {} \n'.format(epoch, accuracy))
valfile.write('Epoch {} Accuracy {} \n'.format(epoch, accuracy))
valfile.close()
datastorage.split(6)
pixel = Pixel('pixel.pkl')
pretrained = PreTrainedResnet({'intermediate_layers':['fc']})
train_dataset = dataset(datastorage, 6, pixel, pretrained, 'train')
train_dataset.create()
val_dataset = dataset(datastorage, 6, pixel, pretrained, 'val')
val_dataset.create()
trackframesbatch, detectionframebatch, labelsbatch, trackcoordsbatch, detectionbatch = val_dataset.collate_fn([val_dataset[0], val_dataset[1]])
print(trackframesbatch)
print(detectionframebatch)
print(labelsbatch)
print(trackcoordsbatch)
print(detectionbatch)'''
test_folder = 'MOT17/train/'
datastorage = datastorage(test_folder)
datastorage.preparetest('DPM')
pixel = Pixel('pixel.pkl')
pretrained = PreTrainedResnet({'intermediate_layers': ['fc']})
test_dataset = testset(datastorage, pixel, pretrained)
test_dataset.create()
detectionsbatch, frameidsbatch, dirtypesbatch, framesbatch = test_dataset.collate_fn(
[test_dataset[0], test_dataset[1]])
print(framesbatch[0].size())
def eval_tracker(use_pickle=False):
trackingmodel = load_model(os.getcwd(), 'tracker.pth', TrackingModel)
trackingmodel.eval()
eval_batch_size = 1
directory_type = 'DPM'
eval_folder = 'MOT17/eval/'
pixel = Pixel('pixel.pkl')
if use_pickle: pixel.load()
pretrained = PreTrainedResnet({'intermediate_layers': ['fc']})
if torch.cuda.is_available():
pretrained = pretrained.cuda()
eval_dataloader, eval_data_size = \
get_eval_data(eval_batch_size, directory_type, pixel, pretrained)
trackingmodel.track(30)
prev_dir = None
tracks = []
def savetracks(tracks, pdir):
file = open(os.path.join(eval_folder, pdir + '.txt'), 'w')
nptracks = np.stack(tracks) #num_frames * num_tracks * 4
nptracks = np.moveaxis(nptracks, 0, 1) #num_tracks * num_frames * 4
num_tracks, num_frames, num_coords = nptracks.shape
for track in range(num_tracks):
for frame in range(num_frames):
if nptracks[track][frame][0] != -1:
file.write(
'{},{},{:.2f},{:.2f},{:.2f},{:.2f},-1,-1,-1\n'.format(
track + 1, frame + 1, nptracks[track][frame][0],
nptracks[track][frame][1],
nptracks[track][frame][2],
nptracks[track][frame][3]))
file.close()
for i, batch in enumerate(eval_dataloader):
detections = Variable(torch.stack(batch[0]))
frameid = batch[1]
dirtype = batch[2][0]
frametensors = Variable(torch.stack(batch[3]))
if torch.cuda.is_available():
detections = detections.cuda()
frametensors = frametensors.cuda()
if prev_dir != dirtype and prev_dir != None:
print(prev_dir)
savetracks(tracks, prev_dir)
tracks = []
detections = detections.squeeze(0)
indices = trackingmodel.score(frametensors, detections)
frametracks = []
for index in indices:
if index != -1:
frametracks.append(detections[index].data.numpy())
else:
frametracks.append(np.asarray([-1, -1, -1, -1]))
tracks.append(np.stack(frametracks))
prev_dir = dirtype
print(prev_dir)
savetracks(tracks, prev_dir)