def find_non_stationary_clusters(args):
if args['use_gpu']:
print("Using CUDA" if torch.cuda.is_available() else "Using CPU")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
else:
device = "cpu"
network_folder_name = args['folder'].split('/')[-1]
tmp = re.search(r"cn(\d+)-", network_folder_name)
n_clusters = int(tmp.group(1))
save_folder = os.path.join(args['dest_root'], network_folder_name)
if os.path.exists(
os.path.join(save_folder, 'cluster_histogram_for_corr.npy')):
print('{} already exists. skipping'.format(
os.path.join(save_folder, 'cluster_histogram_for_corr.npy')))
return
check_mkdir(save_folder)
with open(os.path.join(args['folder'], 'bestval.txt')) as f:
best_val_dict_str = f.read()
bestval = eval(best_val_dict_str.rstrip())
# Network and weight loading
model_config = model_configs.PspnetCityscapesConfig()
net = model_config.init_network(
n_classes=n_clusters,
for_clustering=False,
output_features=False,
use_original_base=args['use_original_base']).to(device)
net.load_state_dict(
torch.load(os.path.join(args['folder'],
bestval['snapshot'] + '.pth'))) # load weights
net.eval()
# copy network file to save location
copyfile(os.path.join(args['folder'], bestval['snapshot'] + '.pth'),
os.path.join(save_folder, 'weights.pth'))
if args['only_copy_weights']:
print('Only copying weights')
return
# Data loading setup
if args['corr_set'] == 'rc':
corr_set_config = data_configs.RobotcarConfig()
elif args['corr_set'] == 'cmu':
corr_set_config = data_configs.CmuConfig()
elif args['corr_set'] == 'both':
corr_set_config1 = data_configs.CmuConfig()
corr_set_config2 = data_configs.RobotcarConfig()
sliding_crop_im = joint_transforms.SlidingCropImageOnly(
713, args['stride_rate'])
input_transform = model_config.input_transform
pre_validation_transform = model_config.pre_validation_transform
if args['corr_set'] == 'both':
corr_set_val1 = correspondences.Correspondences(
corr_set_config1.correspondence_path,
corr_set_config1.correspondence_im_path,
input_size=(713, 713),
input_transform=None,
joint_transform=None,
listfile=corr_set_config1.correspondence_val_list_file)
corr_set_val2 = correspondences.Correspondences(
corr_set_config2.correspondence_path,
corr_set_config2.correspondence_im_path,
input_size=(713, 713),
input_transform=None,
joint_transform=None,
listfile=corr_set_config2.correspondence_val_list_file)
corr_set_val = merged.Merged([corr_set_val1, corr_set_val2])
else:
corr_set_val = correspondences.Correspondences(
corr_set_config.correspondence_path,
corr_set_config.correspondence_im_path,
input_size=(713, 713),
input_transform=None,
joint_transform=None,
listfile=corr_set_config.correspondence_val_list_file)
# Segmentor
segmentor = Segmentor(net, n_clusters, n_slices_per_pass=4)
# save args
open(os.path.join(save_folder,
str(datetime.datetime.now()) + '.txt'),
'w').write(str(args) + '\n\n')
cluster_histogram_for_correspondences = np.zeros((n_clusters, ),
dtype=np.int64)
cluster_histogram_non_correspondences = np.zeros((n_clusters, ),
dtype=np.int64)
for i in range(0, len(corr_set_val), args['step']):
img1, img2, pts1, pts2, _ = corr_set_val[i]
seg1 = segmentor.run_and_save(
img1,
None,
pre_sliding_crop_transform=pre_validation_transform,
input_transform=input_transform,
sliding_crop=sliding_crop_im,
use_gpu=args['use_gpu'])
seg1 = np.array(seg1)
corr_loc_mask = np.zeros(seg1.shape, dtype=np.bool)
valid_inds = (pts1[0, :] >= 0) & (pts1[0, :] < seg1.shape[1]) & (
pts1[1, :] >= 0) & (pts1[1, :] < seg1.shape[0])
pts1 = pts1[:, valid_inds]
for j in range(pts1.shape[1]):
pt = pts1[:, j]
corr_loc_mask[pt[1], pt[0]] = True
cluster_ids_corr = seg1[corr_loc_mask]
hist_tmp, _ = np.histogram(cluster_ids_corr, np.arange(n_clusters + 1))
cluster_histogram_for_correspondences += hist_tmp
cluster_ids_no_corr = seg1[~corr_loc_mask]
hist_tmp, _ = np.histogram(cluster_ids_no_corr,
np.arange(n_clusters + 1))
cluster_histogram_non_correspondences += hist_tmp
if ((i + 1) % 100) < args['step']:
print('{}/{}'.format(i + 1, len(corr_set_val)))
np.save(os.path.join(save_folder, 'cluster_histogram_for_corr.npy'),
cluster_histogram_for_correspondences)
np.save(os.path.join(save_folder, 'cluster_histogram_non_corr.npy'),
cluster_histogram_non_correspondences)
frac = cluster_histogram_for_correspondences / \
(cluster_histogram_for_correspondences +
cluster_histogram_non_correspondences)
stationary_inds = np.argwhere(frac > 0.01)
np.save(os.path.join(save_folder, 'stationary_inds.npy'), stationary_inds)
print('{} stationary clusters out of {}'.format(
len(stationary_inds), len(cluster_histogram_for_correspondences)))
def segment_images_in_folder(network_file, img_folder, save_folder, args):
# get current available device
if args['use_gpu']:
print("Using CUDA" if torch.cuda.is_available() else "Using CPU")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
else:
device = "cpu"
# Network and weight loading
model_config = model_configs.PspnetCityscapesConfig()
if 'n_classes' in args:
print('Initializing model with %d classes' % args['n_classes'])
net = model_config.init_network(
n_classes=args['n_classes'],
for_clustering=False,
output_features=False,
use_original_base=args['use_original_base']).to(device)
else:
net = model_config.init_network().to(device)
print('load model ' + network_file)
state_dict = torch.load(network_file,
map_location=lambda storage, loc: storage)
# needed since we slightly changed the structure of the network in pspnet
state_dict = rename_keys_to_match(state_dict)
net.load_state_dict(state_dict)
net.eval()
# data loading
input_transform = model_config.input_transform
pre_validation_transform = model_config.pre_validation_transform
# make sure crop size and stride same as during training
sliding_crop = joint_transforms.SlidingCropImageOnly(
713, args['sliding_transform_step'])
check_mkdir(save_folder)
t0 = time.time()
# get all file names
filenames_ims = list()
filenames_segs = list()
print('Scanning %s for images to segment.' % img_folder)
for root, subdirs, files in os.walk(img_folder):
filenames = [f for f in files if f.endswith(args['img_ext'])]
if len(filenames) > 0:
print('Found %d images in %s' % (len(filenames), root))
seg_path = root.replace(img_folder, save_folder)
check_mkdir(seg_path)
filenames_ims += [os.path.join(root, f) for f in filenames]
filenames_segs += [
os.path.join(seg_path, f.replace(args['img_ext'], '.png'))
for f in filenames
]
# Create segmentor
if net.n_classes == 19: # This could be the 19 cityscapes classes
segmentor = Segmentor(net,
net.n_classes,
colorize_fcn=cityscapes.colorize_mask,
n_slices_per_pass=args['n_slices_per_pass'])
else:
segmentor = Segmentor(net,
net.n_classes,
colorize_fcn=None,
n_slices_per_pass=args['n_slices_per_pass'])
count = 1
for im_file, save_path in zip(filenames_ims, filenames_segs):
tnow = time.time()
print("[%d/%d (%.1fs/%.1fs)] %s" %
(count, len(filenames_ims), tnow - t0,
(tnow - t0) / count * len(filenames_ims), im_file))
segmentor.run_and_save(
im_file,
save_path,
pre_sliding_crop_transform=pre_validation_transform,
sliding_crop=sliding_crop,
input_transform=input_transform,
skip_if_seg_exists=True,
use_gpu=args['use_gpu'])
count += 1
tend = time.time()
print('Time: %f' % (tend - t0))