def on_batch_end(self, state):
targets = to_numpy(state.input[self.input_key].detach())
outputs = to_numpy(state.output[self.output_key].detach())
self.targets.extend(targets)
self.predictions.extend(outputs)
def draw_predictions_dual(
input: dict,
output: dict,
image_id_key="image_id",
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
class_colors=[
(0, 0, 0), # 0=background
(0, 255, 0), # no damage (or just 'building' for localization) (green)
(255, 255, 0), # minor damage (yellow)
(255, 128, 0), # major damage (red)
(255, 0, 0), # destroyed (red)
],
):
images = []
num_images = len(input[image_id_key])
for i, image_id in enumerate(range(num_images)):
image_pre = rgb_image_from_tensor(input[INPUT_IMAGE_PRE_KEY][i], mean,
std)
image_pre = cv2.cvtColor(image_pre, cv2.COLOR_RGB2BGR)
image_post = rgb_image_from_tensor(input[INPUT_IMAGE_POST_KEY][i],
mean, std)
image_post = cv2.cvtColor(image_post, cv2.COLOR_RGB2BGR)
image_pre_gt = image_pre.copy()
image_post_gt = image_post.copy()
localization_target = to_numpy(input[INPUT_MASK_PRE_KEY][i].squeeze(0))
damage_target = to_numpy(input[INPUT_MASK_POST_KEY][i])
image_pre_gt = overlay_image_and_mask(image_pre_gt,
localization_target,
class_colors)
image_post_gt = overlay_image_and_mask(image_post_gt, damage_target,
class_colors)
localization_predictions = to_numpy(
output[OUTPUT_MASK_PRE_KEY][i].squeeze(0).sigmoid() > 0.5).astype(
np.uint8)
damage_predictions = to_numpy(
output[OUTPUT_MASK_POST_KEY][i]).argmax(axis=0)
image_pre = overlay_image_and_mask(image_pre, localization_predictions,
class_colors)
image_post = overlay_image_and_mask(image_post, damage_predictions,
class_colors)
overlay_gt = np.column_stack([image_pre_gt, image_post_gt])
overlay = np.column_stack([image_pre, image_post])
overlay = np.row_stack([overlay_gt, overlay])
overlay = longest_max_size(overlay, 1024, cv2.INTER_LINEAR)
cv2.putText(overlay, str(image_id), (10, 15), cv2.FONT_HERSHEY_PLAIN,
1, (250, 250, 250))
images.append(overlay)
return images
def visualize_inria_predictions(input: dict,
output: dict,
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225)):
images = []
for image, target, image_id, logits in zip(input['features'],
input['targets'],
input['image_id'],
output['logits']):
image = rgb_image_from_tensor(image, mean, std)
target = to_numpy(target).squeeze(0)
logits = to_numpy(logits).squeeze(0)
overlay = np.zeros_like(image)
true_mask = target > 0
pred_mask = logits > 0
overlay[true_mask & pred_mask] = np.array(
[0, 250, 0], dtype=overlay.dtype
) # Correct predictions (Hits) painted with green
overlay[true_mask & ~pred_mask] = np.array(
[250, 0, 0], dtype=overlay.dtype) # Misses painted with red
overlay[~true_mask & pred_mask] = np.array(
[250, 250, 0],
dtype=overlay.dtype) # False alarm painted with yellow
overlay = cv2.addWeighted(image, 0.5, overlay, 0.5, 0, dtype=cv2.CV_8U)
cv2.putText(overlay, str(image_id), (10, 15), cv2.FONT_HERSHEY_PLAIN,
1, (250, 250, 250))
images.append(overlay)
return images
def on_batch_end(self, state: RunnerState):
y_true = state.input[self.input_key].detach()
y_pred_raw = state.output[self.output_key].detach()
if self.from_regression:
y_pred = regression_to_class(y_pred_raw)
else:
y_pred = torch.argmax(y_pred_raw, dim=1)
y_pred_raw = to_numpy(y_pred_raw)
y_pred = to_numpy(y_pred).astype(int)
y_true = to_numpy(y_true).astype(int)
image_ids = np.array(state.input['image_id'])
if self.ignore_index is not None:
mask = y_true != self.ignore_index
y_pred_raw = y_pred_raw[mask]
y_pred = y_pred[mask]
y_true = y_true[mask]
image_ids = image_ids[mask]
if len(y_true) == 0:
return
negatives = y_true != y_pred
self.image_ids.extend(image_ids[negatives])
self.y_preds_raw.extend(y_pred_raw[negatives])
self.y_preds.extend(y_pred[negatives])
self.y_trues.extend(y_true[negatives])
def draw_binary_segmentation_predictions(input: dict,
output: dict,
image_key='features',
image_id_key='image_id',
targets_key='targets',
outputs_key='logits',
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225)):
images = []
image_id_input = input[image_id_key] if image_id_key is not None else [None] * len(input[image_key])
for image, target, image_id, logits in zip(input[image_key],
input[targets_key],
image_id_input,
output[outputs_key]):
image = rgb_image_from_tensor(image, mean, std)
target = to_numpy(target).squeeze(0)
logits = to_numpy(logits).squeeze(0)
overlay = image.copy()
true_mask = target > 0
pred_mask = logits > 0
overlay[true_mask & pred_mask] = np.array([0, 250, 0],
dtype=overlay.dtype) # Correct predictions (Hits) painted with green
overlay[true_mask & ~pred_mask] = np.array([250, 0, 0], dtype=overlay.dtype) # Misses painted with red
overlay[~true_mask & pred_mask] = np.array([250, 250, 0],
dtype=overlay.dtype) # False alarm painted with yellow
overlay = cv2.addWeighted(image, 0.5, overlay, 0.5, 0, dtype=cv2.CV_8U)
if image_id is not None:
cv2.putText(overlay, str(image_id), (10, 15), cv2.FONT_HERSHEY_PLAIN, 1, (250, 250, 250))
images.append(overlay)
return images
def draw_predictions(
input: dict,
output: dict,
image_id_key="image_id",
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
class_colors=[
(0, 0, 0), # 0=background
(0, 255,
0), # no damage (or just 'building' for localization) (green)
(255, 255, 0), # minor damage (yellow)
(255, 128, 0), # major damage (red)
(255, 0, 0), # destroyed (red)
(127, 127, 127)
],
max_images=32):
images = []
num_images = len(input[image_id_key])
for i in range(num_images):
image_id = input[INPUT_IMAGE_ID_KEY][i]
image_pre = rgb_image_from_tensor(input[INPUT_IMAGE_KEY][i, 0:3, ...],
mean, std)
image_pre = cv2.cvtColor(image_pre, cv2.COLOR_RGB2BGR)
image_post = rgb_image_from_tensor(input[INPUT_IMAGE_KEY][i, 3:6, ...],
mean, std)
image_post = cv2.cvtColor(image_post, cv2.COLOR_RGB2BGR)
image_pre_gt = image_pre.copy()
image_post_gt = image_post.copy()
damage_target = to_numpy(input[INPUT_MASK_KEY][i])
image_pre_gt = overlay_image_and_mask(image_pre_gt, damage_target,
class_colors)
image_post_gt = overlay_image_and_mask(image_post_gt, damage_target,
class_colors)
damage_predictions = to_numpy(output[INPUT_MASK_KEY][i]).argmax(axis=0)
image_pre = overlay_image_and_mask(image_pre, damage_predictions,
class_colors)
image_post = overlay_image_and_mask(image_post, damage_predictions,
class_colors)
overlay_gt = np.column_stack([image_pre_gt, image_post_gt])
overlay = np.column_stack([image_pre, image_post])
overlay = np.row_stack([overlay_gt, overlay])
overlay = longest_max_size(overlay, 1024, cv2.INTER_LINEAR)
cv2.putText(overlay, str(image_id), (10, 15), cv2.FONT_HERSHEY_PLAIN,
1, (250, 250, 250))
images.append(overlay)
if len(images) >= max_images:
break
return images
def on_batch_end(self, state: RunnerState):
outputs = to_numpy(state.output[self.output_key])
targets = to_numpy(state.input[self.input_key])
outputs = np.argmax(outputs, axis=1)
self.outputs.extend(outputs)
self.targets.extend(targets)
def run_dual_inference_on_dataset(model,
dataset,
output_dir,
batch_size=1,
workers=0):
model = model.cuda()
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
model = model.eval()
data_loader = DataLoader(dataset,
batch_size=batch_size,
pin_memory=True,
num_workers=workers)
os.makedirs(output_dir, exist_ok=True)
for batch in tqdm(data_loader):
image_pre = batch[INPUT_IMAGE_PRE_KEY].cuda(non_blocking=True)
image_post = batch[INPUT_IMAGE_POST_KEY].cuda(non_blocking=True)
image_ids = batch[INPUT_IMAGE_ID_KEY]
output = model(image_pre=image_pre, image_post=image_post)
masks_pre = output[OUTPUT_MASK_PRE_KEY]
if masks_pre.size(2) != 1024 or masks_pre.size(3) != 1024:
masks_pre = F.interpolate(masks_pre,
size=(1024, 1024),
mode="bilinear",
align_corners=False)
masks_pre = to_numpy(masks_pre.squeeze(1)).astype(np.float32)
masks_post = output[OUTPUT_MASK_POST_KEY]
if masks_post.size(2) != 1024 or masks_post.size(3) != 1024:
masks_post = F.interpolate(masks_post,
size=(1024, 1024),
mode="bilinear",
align_corners=False)
masks_post = to_numpy(masks_post).astype(np.float32)
for i, image_id in enumerate(image_ids):
_, _, image_uuid = image_id.split("_")
localization_image = masks_pre[i]
damage_image = masks_post[i]
localization_fname = os.path.join(
output_dir, f"test_localization_{image_uuid}_prediction.png")
localization_image = (localization_image > 0.5).astype(np.uint8)
localization_image = colorize_mask(localization_image)
localization_image.save(localization_fname)
damage_fname = os.path.join(
output_dir, f"test_damage_{image_uuid}_prediction.png")
damage_image = np.argmax(damage_image, axis=0).astype(np.uint8)
damage_image = colorize_mask(damage_image)
damage_image.save(damage_fname)
del data_loader
def test_fliplr_image2label():
x = torch.rand((4, 3, 224, 224))
model = GlobalAvgPool2d(flatten=True)
output = tta.fliplr_image2label(model, x)
np.testing.assert_allclose(to_numpy(output),
to_numpy(x.mean(dim=(2, 3))),
atol=1e-6,
rtol=1e-6)
def test_d4_image2mask():
input = torch.rand((4, 3, 224, 224))
model = NoOp()
output = tta.d4_image2mask(model, input)
np.testing.assert_allclose(to_numpy(output),
to_numpy(input),
atol=1e-6,
rtol=1e-6)
def draw_semantic_segmentation_predictions(input: dict,
output: dict,
class_colors,
mode='overlay',
image_key='features',
image_id_key='image_id',
targets_key='targets',
outputs_key='logits',
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225)):
assert mode in {'overlay', 'side-by-side'}
images = []
image_id_input = input[image_id_key] if image_id_key is not None else [
None
] * len(input[image_key])
for image, target, image_id, logits in zip(input[image_key],
input[targets_key],
image_id_input,
output[outputs_key]):
image = rgb_image_from_tensor(image, mean, std)
logits = to_numpy(logits).argmax(axis=0)
target = to_numpy(target)
if mode == 'overlay':
overlay = image.copy()
for class_index, class_color in enumerate(class_colors):
overlay[logits == class_index, :] = class_color
overlay = cv2.addWeighted(image,
0.5,
overlay,
0.5,
0,
dtype=cv2.CV_8U)
if image_id is not None:
cv2.putText(overlay, str(image_id), (10, 15),
cv2.FONT_HERSHEY_PLAIN, 1, (250, 250, 250))
elif mode == 'side-by-side':
true_mask = np.zeros_like(image)
for class_index, class_color in enumerate(class_colors):
true_mask[target == class_index, :] = class_color
pred_mask = np.zeros_like(image)
for class_index, class_color in enumerate(class_colors):
pred_mask[logits == class_index, :] = class_color
overlay = np.hstack((image, true_mask, pred_mask))
else:
raise ValueError(mode)
images.append(overlay)
return images
def test_fliplr_image2mask():
x = torch.rand((4, 3, 224, 224))
model = NoOp()
output = tta.fliplr_image2mask(model, x)
np.testing.assert_allclose(to_numpy(output),
to_numpy(x),
atol=1e-6,
rtol=1e-6)
def on_batch_end(self, state: RunnerState):
outputs = to_numpy(state.output[self.output_key])
targets = to_numpy(state.input[self.input_key])
num_classes = outputs.shape[1]
outputs = [np.eye(num_classes)[y] for y in np.argmax(outputs, axis=1)]
targets = [np.eye(num_classes)[y] for y in targets]
self.outputs.extend(outputs)
self.targets.extend(targets)
def run_dual_inference_on_dataset_oof(model,
dataset,
output_dir,
batch_size=1,
workers=0):
model = model.cuda()
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
model = model.eval()
data_loader = DataLoader(dataset,
batch_size=batch_size,
pin_memory=True,
num_workers=workers)
os.makedirs(output_dir, exist_ok=True)
for batch in tqdm(data_loader):
image_pre = batch[INPUT_IMAGE_PRE_KEY].cuda(non_blocking=True)
image_post = batch[INPUT_IMAGE_POST_KEY].cuda(non_blocking=True)
image_ids = batch[INPUT_IMAGE_ID_KEY]
output = model(image_pre=image_pre, image_post=image_post)
masks_pre = output[OUTPUT_MASK_PRE_KEY]
if masks_pre.size(2) != 1024 or masks_pre.size(3) != 1024:
masks_pre = F.interpolate(masks_pre,
size=(1024, 1024),
mode="bilinear",
align_corners=False)
masks_pre = to_numpy(masks_pre.squeeze(1)).astype(np.float32)
masks_post = output[OUTPUT_MASK_POST_KEY]
if masks_post.size(2) != 1024 or masks_post.size(3) != 1024:
masks_post = F.interpolate(masks_post,
size=(1024, 1024),
mode="bilinear",
align_corners=False)
masks_post = to_numpy(masks_post).astype(np.float32)
for i, image_id in enumerate(image_ids):
localization_image = masks_pre[i]
damage_image = masks_post[i]
localization_fname = os.path.join(
output_dir, fs.change_extension(image_id, ".npy"))
np.save(localization_fname, localization_image)
damage_fname = os.path.join(
output_dir,
fs.change_extension(image_id.replace("_pre", "_post"), ".npy"))
np.save(damage_fname, damage_image)
del data_loader
def get_probabilities(self, state: RunnerState):
probs = state.output[self.output_key].detach().softmax(dim=1)
indexes = state.input[self.sample_index_key]
if probs.size(2) != 1024 or probs.size(3) != 1024:
probs = F.interpolate(probs,
size=(1024, 1024),
mode="bilinear",
align_corners=False)
return to_numpy(probs), to_numpy(indexes)
def test_d4_speed():
df = defaultdict(list)
n = 100
model = resnet34_unet32().cuda().eval()
x = torch.rand((4, 3, 224, 224)).float().cuda()
y1 = tta.d4_image2mask(model, x)
y2 = tta.d4_image_deaugment(model(tta.d4_image_augment(x)))
np.testing.assert_allclose(to_numpy(y1), to_numpy(y2), atol=1e-6, rtol=1e-6)
for deterministic in [False, True]:
for benchmark in [False, True]:
for dtype in [torch.float16, torch.float32]:
torch.cuda.empty_cache()
torch.backends.cuda.deterministic = deterministic
torch.backends.cuda.benchmark = benchmark
model = resnet34_unet32().to(dtype).cuda().eval()
speed_v1 = 0
speed_v2 = 0
for i in range(n):
x = torch.rand((4, 3, 224, 224)).to(dtype).cuda(non_blocking=False)
start = cv2.getTickCount()
y = tta.d4_image2mask(model, x)
v = y.sum().item()
finish = cv2.getTickCount()
speed_v1 += finish - start
np.testing.assert_allclose(v, v, atol=1e-6, rtol=1e-6)
for i in range(n):
x = torch.rand((4, 3, 224, 224)).to(dtype).cuda(non_blocking=False)
start = cv2.getTickCount()
x_a = tta.d4_image_augment(x)
x_a = model(x_a)
y = tta.d4_image_deaugment(x_a)
v = y.sum().item()
finish = cv2.getTickCount()
speed_v2 += finish - start
np.testing.assert_allclose(v, v, atol=1e-6, rtol=1e-6)
df["mode"].append("fp16" if dtype == torch.float16 else "fp32")
df["deterministic"].append(deterministic)
df["benchmark"].append(benchmark)
df["d4_image2mask (ms)"].append(1000.0 * speed_v1 / (cv2.getTickFrequency() * n))
df["d4_augment (ms)"].append(1000.0 * speed_v2 / (cv2.getTickFrequency() * n))
import pandas as pd
df = pd.DataFrame.from_dict(df)
pd.set_option("display.max_columns", None)
pd.set_option("display.max_rows", None)
print(df)
df.to_csv("tta_eval.csv", index=False)
def on_batch_end(self, state: RunnerState):
outputs = to_numpy(state.output[self.output_key])
targets = to_numpy(state.input[self.input_key])
if self.ignore_index is not None:
mask = targets != self.ignore_index
outputs = outputs[mask]
targets = targets[mask]
self.outputs.extend(outputs)
self.targets.extend(targets)
def on_batch_end(self, state):
targets = to_numpy(state.input[self.input_key].detach())
outputs = to_numpy(state.output[self.output_key].detach())
if self.ignore_index is not None:
mask = targets != self.ignore_index
outputs = outputs[mask]
targets = targets[mask]
self.targets.extend(targets)
self.predictions.extend(outputs)
def visualize_canny_predictions(input, output, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
images = []
for image, target, logits in zip(input['features'], input['targets'], output['logits']):
image = rgb_image_from_tensor(image, mean, std)
target = to_numpy(target).squeeze(0)
logits = to_numpy(logits.sigmoid()).squeeze(0)
overlay = np.zeros_like(image)
overlay[logits > 0.5] += np.array([255, 0, 0], dtype=overlay.dtype)
overlay[target > 0] += np.array([0, 255, 0], dtype=overlay.dtype)
overlay = cv2.addWeighted(image, 0.5, overlay, 0.5, 0, dtype=cv2.CV_8U)
images.append(overlay)
return images
def on_batch_end(self, state: RunnerState):
outputs = to_numpy(torch.sigmoid(state.output[self.output_key]))
targets = to_numpy(state.input[self.input_key])
# outputs = 1*(outputs>0.5)
# targets = np.argmax(targets, axis=1)
if self.ignore_index is not None:
mask = targets != self.ignore_index
outputs = outputs[mask]
targets = targets[mask]
self.outputs.extend(outputs)
self.targets.extend(targets)
def average_predictions(predictions: List[pd.DataFrame],
column: str,
method='mean',
min=None,
max=None) -> pd.DataFrame:
preds = []
for p in predictions:
pred = to_numpy(p[column].values.tolist())
preds.append(pred)
preds = np.row_stack(preds)
if min is not None or max is not None:
preds = np.clip(preds, min, max)
if method == 'mean':
y_pred = np.mean(preds, axis=0)
elif method == 'trim_mean':
y_pred = trim_mean(preds, proportiontocut=0.1, axis=0)
elif method == 'median':
y_pred = np.median(preds, axis=0)
else:
raise KeyError(method)
result = pd.DataFrame.from_dict({
'id_code': predictions[0]['image_id'].values,
'diagnosis': y_pred.tolist()
})
return result
def infer_one(model,
mask,
tile_size=(512, 512),
tile_step=(256, 256),
weight='mean'):
image = mask.cpu().numpy()
image = np.moveaxis(image, 0, -1)
with torch.no_grad():
tiler = ImageSlicer((900, 900),
tile_size=tile_size,
tile_step=tile_step,
weight=weight)
tiles = [np.moveaxis(tile, -1, 0) for tile in tiler.split(image)]
merger = CudaTileMerger(tiler.target_shape, 1, tiler.weight)
for tiles_batch, coords_batch in DataLoader(list(
zip(tiles, tiler.crops)),
batch_size=10,
pin_memory=False):
tiles_batch = tiles_batch.float().cuda()
pred_batch = model(tiles_batch)
tiles_batch.cpu().detach()
merger.integrate_batch(pred_batch, coords_batch)
merged_mask = np.moveaxis(to_numpy(merger.merge()), 0, -1)
merged_mask = tiler.crop_to_orignal_size(merged_mask)
m = merged_mask[..., 0].copy()
return m
def test_pseudolabeling_aptos2015_round1(predictions, output_csv):
print('Saving pseudolabels to ', output_csv)
num_models = len(predictions)
ids, x, y_true, y_average = prepare_inference_datasets(
predictions, use_features=False, use_predictions=True)
for i in range(num_models):
print(
fs.id_from_fname(predictions[i]),
cohen_kappa_score(y_true,
regression_to_class(x[:, i]),
weights='quadratic'))
y_round = to_numpy(regression_to_class(x))
y_major = majority_voting(y_round, axis=1)
y_agreement = y_round == np.expand_dims(y_major, -1)
# y_agreement_all = np.all(y_agreement, axis=1)
# y_agreement_all = np.sum(y_agreement, axis=1) >= 16
y_agreement_all = y_major == y_true
print('Agreement', np.mean(y_agreement_all))
print('Distribution', np.bincount(y_major[y_agreement_all]))
y_true[~y_agreement_all] = -100
print(y_round)
df = pd.DataFrame.from_dict({'id_code': ids, 'diagnosis': y_true})
df.to_csv(output_csv, index=None)
def on_batch_end(self, state: RunnerState):
outputs = to_numpy(state.output[self.output_key])
targets = to_numpy(state.input[self.input_key])
if self.ignore_index is not None:
mask = targets != self.ignore_index
outputs = outputs[mask]
targets = targets[mask]
if self.from_regression:
outputs = regression_to_class(outputs)
else:
outputs = np.argmax(outputs, axis=1)
self.outputs.extend(outputs)
self.targets.extend(targets)
def test_tiles_split_merge_cuda():
if not torch.cuda.is_available():
return
class MaxChannelIntensity(nn.Module):
def __init__(self):
super().__init__()
def forward(self, input):
max_channel, _ = torch.max(input, dim=1, keepdim=True)
return max_channel
image = np.random.random((5000, 5000, 3)).astype(np.uint8)
tiler = ImageSlicer(image.shape,
tile_size=(512, 512),
tile_step=(256, 256),
weight='pyramid')
tiles = [tensor_from_rgb_image(tile) for tile in tiler.split(image)]
model = MaxChannelIntensity().eval().cuda()
merger = CudaTileMerger(tiler.target_shape, 1, tiler.weight)
for tiles_batch, coords_batch in DataLoader(list(zip(tiles, tiler.crops)),
batch_size=8,
pin_memory=True):
tiles_batch = tiles_batch.float().cuda()
pred_batch = model(tiles_batch)
merger.integrate_batch(pred_batch, coords_batch)
merged = np.moveaxis(to_numpy(merger.merge()), 0, -1).astype(np.uint8)
merged = tiler.crop_to_orignal_size(merged)
np.testing.assert_equal(merged, image.max(axis=2, keepdims=True))
def on_loader_end(self, state: State):
predictions = to_numpy(self.predictions)
predictions = np.argmax(predictions, axis=1)
score = cohen_kappa_score(predictions,
self.targets,
weights='quadratic')
state.loader_metrics[self.prefix] = score
def predict(model: nn.Module,
image: np.ndarray,
image_size,
normalize=A.Normalize(),
batch_size=1) -> np.ndarray:
tile_step = (image_size[0] // 2, image_size[1] // 2)
tile_slicer = ImageSlicer(image.shape, image_size, tile_step)
tile_merger = CudaTileMerger(tile_slicer.target_shape, 1,
tile_slicer.weight)
patches = tile_slicer.split(image)
transform = A.Compose([normalize, A.Lambda(image=_tensor_from_rgb_image)])
data = list({
"image": patch,
"coords": np.array(coords, dtype=np.int)
} for (patch, coords) in zip(patches, tile_slicer.crops))
for batch in DataLoader(InMemoryDataset(data, transform),
pin_memory=True,
batch_size=batch_size):
image = batch["image"].cuda(non_blocking=True)
coords = batch["coords"]
mask_batch = model(image)
tile_merger.integrate_batch(mask_batch, coords)
mask = tile_merger.merge()
mask = np.moveaxis(to_numpy(mask), 0, -1)
mask = tile_slicer.crop_to_orignal_size(mask)
return mask
def inference_tiles(inference_model,
img_full,
device='cuda',
shape=(32, 1, 768, 448),
weight='mean',
mean=88.904434,
std=62.048634,
plot=False):
bs = shape[0]
input_x = shape[2]
input_y = shape[3]
# Cut large image into overlapping tiles
tiler = ImageSlicer(img_full.shape,
tile_size=(input_x, input_y),
tile_step=(input_x // 2, input_y // 2),
weight=weight)
# HCW -> CHW. Optionally, do normalization here
tiles = [
tensor_from_rgb_image(tile)
for tile in tiler.split(cv2.cvtColor(img_full, cv2.COLOR_GRAY2RGB))
]
# Allocate a CUDA buffer for holding entire mask
merger = CudaTileMerger(tiler.target_shape,
channels=1,
weight=tiler.weight)
# Run predictions for tiles and accumulate them
for tiles_batch, coords_batch in DataLoader(list(zip(tiles, tiler.crops)),
batch_size=bs,
pin_memory=True):
# Move tile to GPU
tiles_batch = ((tiles_batch.float() - mean) / std).to(device)
# Predict
pred_batch = inference_model(tiles_batch)
# Merge on GPU
merger.integrate_batch(pred_batch, coords_batch)
if plot:
for i in range(pred_batch.to('cpu').numpy().shape[0]):
plt.imshow(tiles_batch.to('cpu').numpy()[i, 0, :, :])
plt.show()
plt.imshow(pred_batch.to('cpu').numpy()[i, 0, :, :])
plt.colorbar()
plt.show()
# Normalize accumulated mask and convert back to numpy
merged_mask = np.moveaxis(to_numpy(merger.merge()), 0,
-1).astype('float32')
merged_mask = tiler.crop_to_orignal_size(merged_mask)
torch.cuda.empty_cache()
return merged_mask.squeeze()
def on_batch_end(self, state: RunnerState):
outputs = to_numpy(state.output[self.output_key])
targets = to_numpy(state.input[self.input_key])
num_classes = outputs.shape[1]
outputs = np.argmax(outputs, axis=1)
if self.ignore_index is not None:
mask = targets != self.ignore_index
outputs = outputs[mask]
targets = targets[mask]
outputs = [np.eye(num_classes)[y] for y in outputs]
targets = [np.eye(num_classes)[y] for y in targets]
self.outputs.extend(outputs)
self.targets.extend(targets)
def draw_classification_predictions(input: dict,
output: dict,
class_names,
image_key='image',
image_id_key='image_id',
targets_key='targets',
outputs_key='logits',
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225)):
images = []
for image, target, image_id, logits in zip(input[image_key],
input[targets_key],
input[image_id_key],
output[outputs_key]):
image = rgb_image_from_tensor(image, mean, std)
num_classes = logits.size(0)
target = int(to_numpy(target).squeeze(0))
if num_classes == 1:
logits = int(to_numpy(logits).squeeze(0) > 0)
else:
logits = np.argmax(to_numpy(logits))
overlay = image.copy()
if target != UNLABELED_CLASS:
target_name = class_names[target]
else:
target_name = 'Unlabeled'
cv2.putText(overlay, str(image_id), (10, 15), cv2.FONT_HERSHEY_PLAIN,
1, (250, 250, 250))
cv2.putText(overlay, target_name, (10, 30), cv2.FONT_HERSHEY_PLAIN, 1,
(0, 250, 0))
if target == logits:
cv2.putText(overlay, class_names[logits], (10, 45),
cv2.FONT_HERSHEY_PLAIN, 1, (0, 250, 0))
else:
cv2.putText(overlay, class_names[logits], (10, 45),
cv2.FONT_HERSHEY_PLAIN, 1, (250, 0, 0))
images.append(overlay)
return images
