def test_item_to_mat_conversion():
""" Test conversion mat > item > mat
NOT testing item > mat > item, because comparing items more tedious (first need association!),
"""
# get "original" mat
params = {'model': {'factor_downsample': 4,
'input_height': 512,
'input_width': 1536,
},
}
dataset = data_loader.DataSet(path_csv='../data/train.csv',
path_folder_images='../data/train_images',
path_folder_masks='../data/train_masks',
)
list_ids = dataset.list_ids
item_org = dataset.load_item(list_ids[0], flag_load_mask=False)
dataset_torch = data_loader_torch.DataSetTorch(dataset, params)
mat_org = dataset_torch.convert_item_to_mat(item_org)
# convert back
item = dataset_torch.convert_mat_to_item(mat_org)
item.img = item_org.img # need this for back conversion!
mat = dataset_torch.convert_item_to_mat(item)
assert np.allclose(mat, mat_org)
def evaluate(
model,
device,
params,
):
# define dataset
dataset = data_loader.DataSet(
path_csv=params['datasets']['valid']['path_csv'],
path_folder_images=params['datasets']['valid']['path_folder_images'],
path_folder_masks=params['datasets']['valid']['path_folder_masks'],
)
dataset_torch = data_loader_torch.DataSetTorch(dataset,
params,
flag_augment=False)
dataset_loader = torch.utils.data.DataLoader(
dataset=dataset_torch,
batch_size=params['train']['batch_size_eval'],
shuffle=False,
num_workers=0,
)
# set model to eval (affects e.g. dropout layers) and disable unnecessary grad computation
model.eval()
torch.set_grad_enabled(False)
torch.cuda.empty_cache() # empty cuda cache to prevent memory errors
gc.collect(
) # empty unreferenced objects at the end. Not sure whether better at the end?
# calculate loss for whole dataset
num_batches = len(dataset_loader)
loss_per_name = dict()
print("Evaluating")
for img_batch, mask_batch, heatmap_batch, regr_batch in tqdm(
dataset_loader):
# concat img and mask and perform inference
input = torch.cat([img_batch, mask_batch],
1) # nbatch, nchannels, height, width
output = model(input.to(device))
# calculate loss
batch_loss_per_name = calc_loss(
output,
heatmap_batch.to(device),
regr_batch.to(device),
params['train']['loss'],
)
for name, batch_loss in batch_loss_per_name.items():
if name not in loss_per_name:
loss_per_name[name] = 0
loss_per_name[name] += batch_loss.data
# calculate average
for name, loss in loss_per_name.items():
loss_per_name[name] = loss.cpu().numpy() / len(dataset_loader)
len_dataset = len(dataset_loader.dataset) # check difference
len_dataset2 = len(dataset_loader)
return loss_per_name
def calc_distance(xyz):
# calculate fitting error
x, y, z = xyz
y_pred = 1.0392211185855782 + 0.05107277 * x + 0.16864302 * z # from notebook
dist_fit = (y_pred - y) ** 2
# calculate u,v for reduced image
IMG_SHAPE = (2710, 3384, 3) # img.shape = h,w,c
cam_K = np.array([[2304.5479, 0, 1686.2379],
[0, 2305.8757, 1354.9849],
[0, 0, 1]], dtype=np.float32)
xyz = np.array([x, y, z])
uv = data_loader.xyz2uv(xyz, cam_K)
dataset_torch = data_loader_torch.DataSetTorch(None, params)
uv_new = dataset_torch.convert_uv_to_uv_preprocessed(uv, IMG_SHAPE)
u_new, v_new = uv_new[0], uv_new[1]
dist_uv = (v_new - v_pred) ** 2 + (u_new - u_pred) ** 2
# total
dist_tot = max(0.2, dist_uv) + max(0.4, dist_fit)
# dist_tot = dist_uv + dist_fit
return dist_tot
def predict(model,
device,
params,
):
# Create data generators - they will produce batches
dataset_test = data_loader.DataSet(
path_csv=params['datasets']['test']['path_csv'],
path_folder_images=params['datasets']['test']['path_folder_images'],
path_folder_masks=params['datasets']['test']['path_folder_masks'],
)
dataset_torch_test = data_loader_torch.DataSetTorch(
dataset_test, params,
flag_load_label=False,
flag_augment=False,
)
data_loader_test = torch.utils.data.DataLoader(
dataset=dataset_torch_test,
batch_size=params['predict']['batch_size'],
shuffle=False,
num_workers=4,
pin_memory=True, # see https://pytorch.org/docs/stable/data.html
)
# perform predictions
predictions = []
model.eval()
idx_batch = -1
for img, mask, _, _ in tqdm(data_loader_test):
idx_batch += 1
if idx_batch > params['predict']['num_batches_max']:
print("Ending early because of param num_batches_max={}".format(params['predict']['num_batches_max']))
break
# perform prediction
with torch.no_grad():
# concat img and mask and perform inference
input = torch.cat([img, mask], 1) # nbatch, nchannels, height, width
output = model(input.to(device))
output = output.data.cpu().numpy()
# extract cars as string from each element in batch
num_elems_in_batch = output.shape[0]
for idx_elem_in_batch in range(num_elems_in_batch):
idx_id = idx_batch * params['predict']['batch_size'] + idx_elem_in_batch
id = dataset_test.list_ids[idx_id]
# get mat from output and plot
mat = output[idx_elem_in_batch, ...]
mat = np.rollaxis(mat, 0, 3) # reverse rolling backwards
if params['predict']['flag_plot_mat']:
# convert image to numpy
img_numpy = img.data.cpu().numpy()
img_numpy = img_numpy[idx_elem_in_batch, ...]
img_numpy = np.rollaxis(img_numpy, 0, 3) # reverse rolling backwards
img_numpy = img_numpy[:, :, ::-1] # BGR to RGB
fig, ax = plt.subplots(2, 1, figsize=(10, 10))
ax[0].imshow(img_numpy)
ax_mask = ax[1].imshow(mat[:, :, 0], cmap='PiYG', vmin=-1, vmax=+1)
fig.colorbar(ax_mask, ax=ax[1])
if False: # only use in case of multiple axes, here x,y,z
ax[2].imshow(mat[:, :, 4])
ax[3].imshow(mat[:, :, 5])
ax[4].imshow(mat[:, :, 7])
for axi, label in zip(ax, ['img', 'mask']): # , 'x', 'y', 'z']):
axi.set_ylabel(label)
fig.suptitle('ImageID={}'.format(id))
# save
path_out = os.path.join(params['path_folder_out'],
'pred_mat',
'{:05d}.png'.format(idx_id),
)
os.makedirs(os.path.dirname(path_out), exist_ok=True)
fig.savefig(path_out)
plt.close()
# convert mat to item and plot.
item = dataset_torch_test.convert_mat_to_item(mat)
if params['predict']['flag_optimize']:
for idx_car, car in enumerate(item.cars):
x_new, y_new, z_new, is_marked = optimize_xyz(car.v, car.u, car.x, car.y, car.z, params)
item.cars[idx_car].x = x_new
item.cars[idx_car].y = y_new
item.cars[idx_car].z = z_new
item.cars[idx_car].is_marked = is_marked
if params['predict']['flag_plot_item']:
item_org = dataset_test.load_item(id)
item.img = item_org.img
item.mask = np.zeros((1, 1))
fig, ax = item.plot()
fig.suptitle('ImageID={}'.format(id))
if idx_batch == 2:
num_cars = len(item.cars)
plt.show()
# save
path_out = os.path.join(params['path_folder_out'],
'pred_item',
'{:05d}.png'.format(idx_id),
)
os.makedirs(os.path.dirname(path_out), exist_ok=True)
fig.savefig(path_out)
plt.close()
# extract prediction string from item
string = item.get_cars_as_string(flag_submission=True)
predictions.append(string)
# predictions to csv
df_out = pd.DataFrame()
df_out['ImageId'] = dataset_test.list_ids
df_out.loc[0:len(predictions) - 1, 'PredictionString'] = predictions
print(df_out.head())
path_csv = os.path.join(params['path_folder_out'], 'predictions.csv')
df_out.to_csv(path_csv, sep=',', index=False)
return df_out
def train(
model,
device,
params,
):
# define training dataset
dataset = data_loader.DataSet(
path_csv=params['datasets']['train']['path_csv'],
path_folder_images=params['datasets']['train']['path_folder_images'],
path_folder_masks=params['datasets']['train']['path_folder_masks'],
)
dataset_torch = data_loader_torch.DataSetTorch(
dataset,
params,
flag_load_label=True,
flag_augment=params['train']['use_augmentation'],
)
dataset_loader = torch.utils.data.DataLoader(
dataset=dataset_torch,
batch_size=params['train']['batch_size'],
shuffle=True,
num_workers=4,
pin_memory=True, # see https://pytorch.org/docs/stable/data.html
)
# define optimizer and decrease learning rate by 0.1 every 3 epochs
optimizer = torch.optim.AdamW(
model.parameters(),
lr=params['train']['learning_rate']['initial'],
)
lr_scheduler = torch.optim.lr_scheduler.StepLR(
optimizer,
step_size=params['train']['learning_rate']['num_epochs_const'],
gamma=params['train']['learning_rate']['factor_decrease'],
)
# for each epoch...
df_out = pd.DataFrame()
for idx_epoch in range(params['train']['num_epochs']):
print("Training epoch {}".format(idx_epoch))
# set model to train (affects e.g. dropout layers) and disable unnecessary grad computation
model.train()
torch.set_grad_enabled(True)
torch.cuda.empty_cache() # empty cuda cache to prevent memory errors
gc.collect(
) # empty unreferenced objects at the end. Not sure whether better at the end?
# calculate loss for whole dataset
num_batches = len(dataset_loader)
loss_per_name = dict()
dataset_tqdm = tqdm(dataset_loader)
for img_batch, mask_batch, heatmap_batch, regr_batch in dataset_tqdm:
# concat img and mask and perform inference
input = torch.cat([img_batch, mask_batch],
1) # nbatch, nchannels, height, width
output = model(input.to(device))
# calculate loss
batch_loss_per_name = calc_loss(
output,
heatmap_batch.to(device),
regr_batch.to(device),
params['train']['loss'],
)
for name, batch_loss in batch_loss_per_name.items():
if name not in loss_per_name:
loss_per_name[name] = 0
loss_per_name[name] += batch_loss.data
# change tqdm progress bar description
description = "loss: "
for name, batch_loss in batch_loss_per_name.items():
description += "{}={:.3f} ".format(
name,
batch_loss.data.cpu().numpy())
dataset_tqdm.set_description(description)
# perform optimization
batch_loss_per_name['tot'].backward(
) # computes x.grad += dloss/dx for all parameters x
optimizer.step() # updates values x += -lr * x.grad
optimizer.zero_grad() # set x.grad = 0, for next iteration
# step learning rate after each epoch (not after each batch)
lr_scheduler.step()
# calculate average and store results
for name in loss_per_name.keys():
loss_per_name[name] = loss_per_name[name].cpu().numpy() / len(
dataset_loader)
df_out.loc[idx_epoch, 'loss_' + name] = loss_per_name[name]
values_per_name = evaluate(model, device, params)
for key, value in values_per_name.items():
df_out.loc[idx_epoch, 'valid_' + key] = value
# save history
path_csv = os.path.join(
params['path_folder_out'],
'train_history_{}.csv'.format(idx_epoch),
)
os.makedirs(os.path.dirname(path_csv), exist_ok=True)
df_out.to_csv(path_csv, sep=';')
print(df_out)
# save model weights
path_weights = os.path.join(
params['path_folder_out'],
'model_{}.pth'.format(idx_epoch),
)
torch.save(model.state_dict(), path_weights)
return df_out