File size: 2,424 Bytes
8966f37 | 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 | from typing import Optional
import torch
from transformers import PreTrainedModel
from transformers.modeling_outputs import SemanticSegmenterOutput
from .configuration_cond_unet import CondUNetConfig
try:
from .unet_attn import UNet2DAttn
except ModuleNotFoundError:
# The release exporter bundles these modules into the Hub repository root.
from nets.unet_attn import UNet2DAttn
class CondUNetForSemanticSegmentation(PreTrainedModel):
config_class = CondUNetConfig
main_input_name = "pixel_values"
def __init__(self, config: CondUNetConfig):
super().__init__(config)
self.unet = UNet2DAttn(
in_channels=config.in_channels,
num_classes=config.num_labels,
n_organs=config.n_organs,
size=config.size,
depth=config.depth,
attn_start=config.attn_start,
use_attn=config.use_attn,
img_size=config.image_size,
patch_size=config.patch_size,
emb_dim=config.emb_dim,
n_heads=config.n_heads,
distill=False,
distill_unet=False,
use_dwt=config.use_dwt,
wavelet=config.wavelet,
dwt_bands=config.dwt_bands,
use_shape=config.use_shape,
shape_res=config.shape_res,
)
self.post_init()
def forward(
self,
pixel_values: torch.FloatTensor,
organ_id: Optional[torch.LongTensor] = None,
labels: Optional[torch.FloatTensor] = None,
return_dict: Optional[bool] = None,
**kwargs,
):
if organ_id is None:
organ_id = torch.full(
(pixel_values.shape[0],),
self.config.unknown_organ_id,
device=pixel_values.device,
dtype=torch.long,
)
else:
organ_id = organ_id.to(device=pixel_values.device, dtype=torch.long)
outputs = self.unet(
pixel_values=pixel_values,
organ_id=organ_id,
masks=labels,
**kwargs,
)
logits = outputs["logits"]
if logits.ndim == 3:
logits = logits.unsqueeze(1)
loss = outputs["loss"] if labels is not None else None
if return_dict is False:
return (loss, logits) if loss is not None else (logits,)
return SemanticSegmenterOutput(loss=loss, logits=logits)
|