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 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_semseg( dataset_root, model_name, batch_size=8, image_size=256, use_gt_labels=False, ): print('Save semseg : ', locals()) import shutil import cv2 import os import numpy as np import torch.nn.functional as F import torch from models.get_model import get_model from torch.utils.data.dataloader import DataLoader from data.loader_segmentation import Segmentation from artifacts.artifact_manager import artifact_manager from data.voc2012 import label_to_image # Set up model model = get_model(model_name) model.load() model.to(model.device) model.train(False) # Set up data loader dataloader = DataLoader( Segmentation(dataset_root, source='val', source_augmentation='val', image_size=image_size, requested_labels=['classification', 'segmentation']), batch_size=batch_size, shuffle=False, num_workers=4, prefetch_factor=4, ) # Clear and create destination directory semseg_path = os.path.join(artifact_manager.getDir(), 'semseg_output') if (os.path.exists(semseg_path)): shutil.rmtree(semseg_path) os.makedirs(semseg_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(): semsegs = model.event({ 'name': 'get_semseg', 'inputs': inputs_in, 'labels': labels_in, 'batch': batch_no + 1 }) semsegs = semsegs.detach().cpu().numpy() # Save out cams for semseg_no, semseg in enumerate(semsegs): # Save out ground truth labels for testing the rest of the system if use_gt_labels: semseg = labels_in['segmentation'][semseg_no][1:] semseg = F.adaptive_avg_pool2d(semseg, [32, 32]).numpy() for i in range(0, semseg.shape[0]): semseg[i] = cv2.blur(semseg[i], (3, 3)) semseg[i] = cv2.blur(semseg[i], (3, 3)) # # Disregard false positives # gt_mask = labels_in['classification'][semseg_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 # Upsample CAM to original image size # - Calculate original image aspect ratio width = datapacket_in['width'][semseg_no].detach().numpy() height = datapacket_in['height'][semseg_no].detach().numpy() aspect_ratio = width / height # - Calculate width and height to cut from upscaled CAM if aspect_ratio > 1: cut_width = image_size cut_height = round(image_size / aspect_ratio) else: cut_width = round(image_size * aspect_ratio) cut_height = image_size # - Upscale CAM to match input size semseg = np.moveaxis(semseg, 0, -1) semseg = cv2.resize(semseg, (image_size, image_size), interpolation=cv2.INTER_LINEAR) semseg = np.moveaxis(semseg, -1, 0) # - Cut CAM from input size and upscale to original image size semseg = semseg[:, 0:cut_height, 0:cut_width] semseg = np.moveaxis(semseg, 0, -1) semseg = cv2.resize(semseg, (width, height), interpolation=cv2.INTER_LINEAR) semseg = np.moveaxis(semseg, -1, 0) semseg_as_label = label_to_image(semseg) # Write image img_no = datapacket_in['image_name'][semseg_no] cv2.imwrite( os.path.join(semseg_path, img_no) + '.png', semseg_as_label * 255) print('Save cam : ', img_no, end='\r') print('')
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_labels( dataset_root, model_name, batch_size=8, image_size=256, use_gt_labels=False, ): print('Save cams : ', locals()) 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(model_name) model.load() model.to(model.device) # Set up data loader dataloader = DataLoader( Segmentation( dataset_root, source='trainval', source_augmentation='val', image_size=image_size, requested_labels=['classification', 'segmentation'] ), batch_size=batch_size, shuffle=False, num_workers=4, ) # Clear and create desintation directory cam_path = os.path.join(artifact_manager.getDir(), 'cam') if (os.path.exists(cam_path)): shutil.rmtree(cam_path) os.makedirs(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 use_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 # Upsample CAM to original image size # - Calculate original image aspect ratio width = datapacket_in['width'][cam_no].detach().numpy() height = datapacket_in['height'][cam_no].detach().numpy() aspect_ratio = width / height # - Calculate width and height to cut from upscaled CAM if aspect_ratio > 1: cut_width = image_size cut_height = round(image_size / aspect_ratio) else: cut_width = round(image_size * aspect_ratio) cut_height = image_size # - Upscale CAM to match input size cam = np.moveaxis(cam, 0, -1) cam = cv2.resize(cam, (image_size, image_size), interpolation=cv2.INTER_LINEAR) cam = np.moveaxis(cam, -1, 0) # - Cut CAM from input size and upscale to original image size cam = cam[:, 0:cut_height, 0:cut_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) # 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 # Write image img_no = datapacket_in['image_name'][cam_no] cv2.imwrite(os.path.join(cam_path, img_no) + '.png', cam_norm * 255) print('Save cam : ', img_no, 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('')