File size: 3,377 Bytes
74e8f2f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
# Copyright 2022 Big Vision Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Inputs, outputs and losses for panoptic task."""
import big_vision.utils as u
import einops
import jax
import jax.numpy as jnp
import numpy as np
ONE_HOT_AXIS = -2
def input_pp(batch, config):
"""Make inputs for panoptic segmentation task."""
if "labels" not in batch:
# During predict of phase2 there is no 'labels' field.
x = None
else:
hp, wp = config.model.patch_size
x = {
"semantics": batch["labels"][..., 0],
"instances": batch["labels"][..., 1],
}
# Convert labels from (B, H, W) to (B, num_patches, num_classes, patch_size)
for key in ["semantics", "instances"]:
x[key] = jax.nn.one_hot(
einops.rearrange(
x[key], "b (hn hp) (wn wp) -> b (hn wn) (hp wp)", hp=hp, wp=wp),
num_classes=config.model.inputs[key][ONE_HOT_AXIS], axis=ONE_HOT_AXIS)
ctx = batch.get("image_ctx", batch.get("image", None))
return {"ctx": ctx, "x": x}
def loss_fn(logits, batch, config):
"""Compute loss for panoptic task."""
labels = input_pp(batch, config)["x"]
losses = {}
for key in ["semantics", "instances"]:
losses[f"loss_{key}"] = u.softmax_xent(
logits=logits[key], labels=labels[key], reduction=False,
axis=ONE_HOT_AXIS)
return sum(losses.values()), losses
def predict_outputs(logits, config, min_fraction=0.0):
"""Make outputs for panoptic segmentation task."""
# Map logits to (height, width, channels).
hp, wp = config.model.patch_size
hn, wn = np.array(config.model.input_size) // np.array((hp, wp))
outputs = {}
for key in ["semantics", "instances"]:
assert ONE_HOT_AXIS == -2, "Rearrange below depends on this."
outputs[key] = einops.rearrange(
logits[key],
"b (hn wn) c (hp wp) -> b (hn hp) (wn wp) c",
hn=hn, wn=wn, hp=hp, wp=wp)
return panoptic_predictions_from_logits(
**outputs, min_fraction=min_fraction)
def panoptic_predictions_from_logits(semantics, instances, min_fraction=0.0):
"""Make panoptic prediction from logits."""
ins = jnp.argmax(instances, axis=-1)
# Note: Make sure each instance has all pixels annotated with same label.
# Otherwise they are further split into more instances and greatly affect
# the number of unmatched predicted segments (FP) and RQ.
masks = jax.nn.one_hot(ins, instances.shape[-1], dtype=jnp.int32)
label = jnp.argmax(jnp.einsum("bhwk,bhwn->bnk", semantics, masks), axis=-1)
sem = jnp.einsum("bhwn,bn->bhw", masks, label)
out = jnp.stack([sem, ins], axis=-1)
# Filter out small objects
fraction = jnp.sum(masks, axis=(1, 2), keepdims=True)/np.prod(ins.shape[1:3])
mask_big = (fraction > min_fraction).astype("int32")
mask_big_spatial = jnp.sum(masks * mask_big, axis=-1, keepdims=True) > 0
return out * mask_big_spatial.astype("int32")
|