def train_classifier(config: Config):
config_json = config.toDictionary()
print('train_classifier')
print(config_json)
from training.train import train
from torch.utils.data.dataloader import DataLoader
from data.loader_segmentation import Segmentation
model = get_model(config.classifier_name)
wandb.init(entity='kobus_wits', project='wass_classifier', name=config.sweep_id + '_c_' + config.classifier_name, config=config_json)
wandb.watch(model)
train(
model=model,
dataloaders = {
'train': DataLoader(
Segmentation(
config.classifier_dataset_root,
source='train',
augmentation='train',
image_size=config.classifier_image_size
),
batch_size=config.classifier_batch_size_train,
shuffle=True,
pin_memory=True,
num_workers=4,
prefetch_factor=4
),
},
epochs=config.classifier_epochs,
validation_mod=10
)
wandb.finish()
def start():
from time import time
from training.train_classifier import train_classifier
from training.train_affinitynet import train_affinitynet
from training.save_cams import save_cams
from training.save_cams import measure_cams
from training.save_cams_crf import save_cams_crf
from training.save_cams_random_walk import save_cams_random_walk
from training.save_cams_random_walk import measure_random_walk
from training.train_semseg import train_semseg
from training.save_semseg import save_semseg
from training.save_semseg import measure_semseg
from artifacts.artifact_manager import artifact_manager
from training.config_manager import Config
from training.config_manager import config_manager
sweeps = [
Config(
eval_dataset_root='datasets/generated/voc',
classifier_dataset_root='datasets/generated/voc_aug',
classifier_name='vgg16',
classifier_epochs=1,
classifier_batch_size_train=32,
classifier_pretrained=True,
classifier_pretrained_unfreeze=10,
cams_produce_batch_size=32,
cams_measure_batch_size=64,
affinity_net_batch_size=8,
),
]
for sweep_index, sweep in enumerate(sweeps):
artifact_manager.setArtifactContainer('psa_sweep_' + str(sweep_index))
sweep.sweep_id = str(time())
config_manager.setConfig(sweep)
print(f'Sweep start {sweep_index}/{len(sweeps)}')
# ########## CAMS
# # Train the classifier
# train_classifier(sweep)
# # Save out the CAMs
# save_cams(sweep)
# # Measure cams
# measure_cams(sweep)
# ########## AFFINITY NET
# # Apply DCRF on CAMs
# save_cams_crf(sweep)
# # Train AffinityNet
train_affinitynet(sweep)
def train_affinitynet(config: Config):
config_json = config.toDictionary()
print('train_affinitynet')
print(config_json)
from training.train import train
from torch.utils.data.dataloader import DataLoader
from data.loader_segmentation import Segmentation
from artifacts.artifact_manager import artifact_manager
model = get_model(config.affinity_net_name)
wandb.init(entity='kobus_wits', project='wass_affinity', name=config.sweep_id + '_a_' + config.affinity_net_name, config=config_json)
wandb.watch(model)
train(
model=model,
dataloaders = {
'train': DataLoader(
Segmentation(
config.classifier_dataset_root,
source='train',
augmentation='train',
image_size=config.affinity_net_image_size,
requested_labels=['affinity'],
affinity_root=artifact_manager.getDir()
),
batch_size=config.affinity_net_batch_size,
shuffle=False,
pin_memory=False,
num_workers=4,
prefetch_factor=4
),
},
epochs=config.affinity_net_epochs,
validation_mod=10
)
wandb.finish()
def save_cams_crf(config: Config):
config_json = config.toDictionary()
print('save_cams_crf')
print(config_json)
import shutil
import numpy as np
from torch.utils.data.dataloader import DataLoader
from data.loader_segmentation import Segmentation
from data.voc2012 import label_to_image
from multiprocessing import Pool
from artifacts.artifact_manager import artifact_manager
cam_root_path = os.path.join(artifact_manager.getDir(), 'cam')
# Set up data loader
dataloader = DataLoader(Segmentation(
config.classifier_dataset_root,
source='train',
image_size=config.classifier_image_size),
batch_size=32,
shuffle=False,
num_workers=2,
prefetch_factor=2)
# Create high, low dirs
cam_la_path = os.path.join(artifact_manager.getDir(), 'cam_la')
if (os.path.exists(cam_la_path)):
shutil.rmtree(cam_la_path)
os.makedirs(cam_la_path)
cam_ha_path = os.path.join(artifact_manager.getDir(), 'cam_ha')
if (os.path.exists(cam_ha_path)):
shutil.rmtree(cam_ha_path)
os.makedirs(cam_ha_path)
wandb.init(entity='kobus_wits',
project='wass_measure_cams_crfs',
name=config.sweep_id + '_cam_' + config.classifier_name,
config=config_json)
count = 0
for batch_no, batch in enumerate(dataloader):
from training.save_cams_crf import _process_sample
labels = batch[1]
datapacket = batch[2]
payloads = []
for image_no, image_name in enumerate(datapacket['image_name']):
payload = {
'image_name': image_name,
'count': count,
'image_width': datapacket['width'][image_no].numpy(),
'image_height': datapacket['height'][image_no].numpy(),
'channels': labels['classification'].shape[1],
'image_path': datapacket['image_path'][image_no],
'cam_la_path': cam_la_path,
'cam_ha_path': cam_ha_path,
'cam_root_path': cam_root_path,
'alpha_low': config.cams_bg_alpha_low,
'alpha_high': config.cams_bg_alpha_high,
}
payloads.append(payload)
count += 1
print('Save cam : ', count, end='\r')
with Pool(8) as poel:
logs = poel.map(_process_sample, payloads)
for log in logs:
wandb.log(log, step=log['image_count'])
print('')
wandb.finish()
def measure_random_walk(config: Config):
config_json = config.toDictionary()
print('measure_cams')
print(config_json)
import os
from torch.utils.data.dataloader import DataLoader
from data.loader_segmentation import Segmentation
from artifacts.artifact_manager import artifact_manager
from multiprocessing import Pool
# Set up data loader
dataloader = DataLoader(Segmentation(
config.eval_dataset_root,
source='train',
augmentation='val',
image_size=config.classifier_image_size,
requested_labels=['classification', 'segmentation']),
batch_size=config.affinity_net_batch_size,
shuffle=False,
num_workers=2,
prefetch_factor=2)
# Get cams directory
labels_rw_root_path = os.path.join(artifact_manager.getDir(), 'labels_rw')
count = 0
wandb.init(entity='kobus_wits',
project='wass_measure_cams_rw',
name=config.sweep_id + '_cam_' + config.classifier_name,
config=config_json)
avg_meter = AverageMeter('accuracy', 'mapr', 'miou')
for batch_no, batch in enumerate(dataloader):
datapacket_in = batch[2]
payloads = []
logs = []
for image_no, image_name in enumerate(datapacket_in['image_name']):
payload = {
'count':
count,
'label_path':
datapacket_in['label_path'][image_no],
'predi_path':
os.path.join(labels_rw_root_path, image_name + '.png'),
}
payloads.append(payload)
logs.append(_measure_sample(payload))
count += 1
print('Measure cam RW : ', count, end='\r')
# with Pool(8) as poel:
# logs = poel.map(_measure_sample, payloads)
for log in logs:
avg_meter.add({
'accuracy': log['accuracy'],
'mapr': log['mapr'],
'miou': log['miou'],
})
if log['count'] < 8:
wandb.log(log, step=log['count'])
wandb.log({
'accuracy': avg_meter.get('accuracy'),
'mapr': avg_meter.get('mapr'),
'miou': avg_meter.get('miou'),
})
wandb.finish()
def save_cams_random_walk(config: Config):
config_json = config.toDictionary()
print('save_cams_random_walk')
print(config_json)
import shutil
import os
from torch.utils.data.dataloader import DataLoader
from data.loader_segmentation import Segmentation
from artifacts.artifact_manager import artifact_manager
# Set up model
model = get_model(config.affinity_net_name)
model.load()
model.eval()
model.to(model.device)
# Set up data loader
dataloader = DataLoader(Segmentation(
config.classifier_dataset_root,
source='train',
augmentation='affinity_predict',
image_size=config.affinity_net_image_size,
requested_labels=['classification', 'segmentation']),
batch_size=1,
shuffle=False,
num_workers=2,
prefetch_factor=2)
# Get cam source directory
cam_path = os.path.join(artifact_manager.getDir(), 'cam')
# Clear and create output directory
labels_rw_path = os.path.join(artifact_manager.getDir(), 'labels_rw')
if (os.path.exists(labels_rw_path)):
shutil.rmtree(labels_rw_path)
os.makedirs(labels_rw_path)
count = 0
for batch_no, batch in enumerate(dataloader):
inputs = batch[0]
labels = batch[1]
datapacket = batch[2]
for image_no, image_name in enumerate(datapacket['image_name']):
image = inputs['image'].cuda(non_blocking=True)
image_width = datapacket['width'][image_no].numpy()
image_height = datapacket['height'][image_no].numpy()
channels = labels['classification'].shape[1]
# Pad image
image_width_padded = int(np.ceil(image_width / 8) * 8)
image_height_padded = int(np.ceil(image_height / 8) * 8)
image_padded = F.pad(image,
(0, image_width_padded - image_width, 0,
image_height_padded - image_height))
image_width_pooled = int(np.ceil(image_width_padded / 8))
image_height_pooled = int(np.ceil(image_height_padded / 8))
# Load cam
cam_path_instance = os.path.join(cam_path, image_name + '.png')
cam = cv2.imread(cam_path_instance, cv2.IMREAD_GRAYSCALE)
cam = np.reshape(cam, ((channels, image_height, image_width)))
cam = cam / 255.0
# Build cam background
cam_background = (
1 - np.max(cam,
(0), keepdims=True))**config.affinity_net_bg_alpha
cam = np.concatenate((cam_background, cam), axis=0)
cam = cam.astype(np.float32)
# Pad cam
cam_padded_width = int(np.ceil(cam.shape[2] / 8) * 8)
cam_padded_height = int(np.ceil(cam.shape[1] / 8) * 8)
cam_padded = np.pad(cam,
((0, 0), (0, cam_padded_height - image_height),
(0, cam_padded_width - image_width)),
mode='constant')
# Run images through model and get affinity matrix
with torch.no_grad():
aff_mat = model.event({
'name': 'infer_aff_net_dense',
'image': image_padded,
})
aff_mat = torch.pow(aff_mat, config.affinity_net_beta)
trans_mat = aff_mat / torch.sum(aff_mat, dim=0, keepdim=True)
for _ in range(config.affinity_net_log_t):
trans_mat = torch.matmul(trans_mat, trans_mat)
cam_pooled = F.avg_pool2d(torch.from_numpy(cam_padded), 8, 8)
cam_vec = cam_pooled.view(21, -1)
cam_rw = torch.matmul(cam_vec.cuda(), trans_mat)
cam_rw = cam_rw.view(1, 21, image_height_pooled,
image_width_pooled)
cam_rw = torch.nn.Upsample(
(image_height_padded, image_width_padded),
mode='bilinear')(cam_rw)
cam_rw = cam_rw.cpu().data[0, :, :image_height, :image_width]
label_rw = label_to_image(cam_rw)
cv2.imwrite(os.path.join(labels_rw_path, image_name + '.png'),
label_rw * 255)
count += 1
print('Save cam : ', count, end='\r')
print('')
def save_cams(config: Config):
config_json = config.toDictionary()
print('save_cams')
print(config_json)
import shutil
import cv2
import os
import numpy as np
from torch.utils.data.dataloader import DataLoader
from data.loader_segmentation import Segmentation
from artifacts.artifact_manager import artifact_manager
# Set up model
model = get_model(config.classifier_name)
model.load()
model.eval()
model.to(model.device)
# Set up data loader
dataloader = DataLoader(Segmentation(
config.classifier_dataset_root,
source='train',
augmentation='val',
image_size=config.classifier_image_size,
requested_labels=['classification', 'segmentation']),
batch_size=config.cams_produce_batch_size,
shuffle=False,
num_workers=4,
prefetch_factor=4)
# Clear and create destination directory
cam_path = os.path.join(artifact_manager.getDir(), 'cam')
if (os.path.exists(cam_path)):
shutil.rmtree(cam_path)
os.makedirs(cam_path)
label_cam_path = os.path.join(artifact_manager.getDir(), 'labels_cam')
if (os.path.exists(label_cam_path)):
shutil.rmtree(label_cam_path)
os.makedirs(label_cam_path)
for batch_no, batch in enumerate(dataloader):
inputs_in = batch[0]
labels_in = batch[1]
datapacket_in = batch[2]
# Run images through model and get raw cams
with torch.no_grad():
cams = model.event({
'name': 'get_cam',
'inputs': inputs_in,
'labels': labels_in,
'batch': batch_no + 1
})
# Save out cams
for cam_no, cam in enumerate(cams):
# Save out ground truth labels for testing the rest of the system
if config.cams_save_gt_labels:
cam = labels_in['segmentation'][cam_no][1:]
cam = F.adaptive_avg_pool2d(cam, [32, 32]).numpy()
for i in range(0, cam.shape[0]):
cam[i] = cv2.blur(cam[i], (3, 3))
cam[i] = cv2.blur(cam[i], (3, 3))
# Disregard false positives
gt_mask = labels_in['classification'][cam_no].numpy()
gt_mask[gt_mask > 0.5] = 1
gt_mask[gt_mask <= 0.5] = 0
gt_mask = np.expand_dims(np.expand_dims(gt_mask, -1), -1)
cam *= gt_mask
# Scale CAM to match input size
cam = np.moveaxis(cam, 0, -1)
cam = cv2.resize(
cam,
(config.classifier_image_size, config.classifier_image_size),
interpolation=cv2.INTER_LINEAR)
cam = np.moveaxis(cam, -1, 0)
# - Cut CAM from input size and upscale to original image size
width = datapacket_in['width'][cam_no].detach().numpy()
height = datapacket_in['height'][cam_no].detach().numpy()
content_width = datapacket_in['content_width'][cam_no].detach(
).numpy()
content_height = datapacket_in['content_height'][cam_no].detach(
).numpy()
cam = cam[:, 0:content_height, 0:content_width]
cam = np.moveaxis(cam, 0, -1)
cam = cv2.resize(cam, (width, height),
interpolation=cv2.INTER_LINEAR)
cam = np.moveaxis(cam, -1, 0)
# Normalize each cam map to between 0 and 1
cam_max = np.max(cam, (1, 2), keepdims=True)
cam_norm = cam / (cam_max + 1e-5)
cam_bg = (
1 -
np.max(cam_norm, axis=0, keepdims=True))**config.cams_bg_alpha
cam_with_bg = np.concatenate((cam_bg, cam_norm), axis=0)
label_cam = label_to_image(cam_with_bg)
# Collapse cam from 3d into long 2d
cam_norm = np.reshape(
cam_norm,
(cam_norm.shape[0] * cam_norm.shape[1], cam_norm.shape[2]))
cam_norm[cam_norm > 1] = 1
cam_norm[cam_norm < 0] = 0
label_cam[label_cam > 1] = 1
label_cam[label_cam < 0] = 0
# Write image
img_no = datapacket_in['image_name'][cam_no]
cv2.imwrite(
os.path.join(cam_path, img_no) + '.png', cam_norm * 255)
cv2.imwrite(
os.path.join(label_cam_path, img_no) + '.png', label_cam * 255)
print('Save cam : ', img_no, end='\r')
print('')