| import math |
| from types import SimpleNamespace |
| import numpy as np |
| import torch |
| from torch import nn |
| from timm.models.vision_transformer import Block |
| from .radio_patch import ViTPatchGenerator |
| from .radio_cls import ClsToken as _bundle_radio_cls |
|
|
| class RadioBackbone(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.patch_generator = ViTPatchGenerator(16,1280,224,cls_token=True,max_input_dims=2048,num_cls_tokens=4,register_multiple=10) |
| self.blocks = nn.Sequential(*[Block(1280,16,mlp_ratio=4,qkv_bias=True) for _ in range(32)]) |
| def forward(self, pixels): |
| return self.blocks(self.patch_generator(pixels))[:,10:] |
|
|
| class RadioModel(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.model = RadioBackbone() |
| self.register_buffer('summary_idxs',torch.tensor([0,1,2],dtype=torch.int64)) |
| def forward(self, pixels): |
| return self.model(pixels) |
|
|
| class RadioWrapper(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.radio_model = RadioModel() |
| def forward(self, pixels): |
| return self.radio_model(pixels) |
|
|
| class RadioVisionTower(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.vision_tower = RadioWrapper() |
| self.image_processor = RadioImageProcessor() |
| self.is_loaded = True |
| self.config = SimpleNamespace(hidden_size=1280,patch_size=16) |
| def forward(self, images): |
| reference = next(self.parameters()) |
| if isinstance(images, (tuple,list)): |
| return torch.cat([self.vision_tower(x.to(reference)) for x in images], dim=1) |
| return self.vision_tower(images.to(reference)) |
|
|
| class RadioImageProcessor: |
| radio_downsample_ratio = 0.5 |
| def __call__(self, images, return_tensors='pt'): |
| if len(images)!=1: |
| raise ValueError('Process one image at a time') |
| image = images[0].convert('RGB') |
| ph,pw = round(image.height/16+0.5), round(image.width/16+0.5) |
| factor=min(math.sqrt(12000/(ph*pw)),1.) |
| h,w=math.floor(factor*ph),math.floor(factor*pw) |
| if h*w<1024: |
| up=math.sqrt(1024/(h*w)); h,w=math.ceil(up*h),math.ceil(up*w) |
| if h%2:h=h+1 if (h+1)*w<=12000 else max(2,h-1) |
| if w%2:w=w+1 if h*(w+1)<=12000 else max(2,w-1) |
| pixels=torch.from_numpy(np.asarray(image,dtype=np.uint8).copy()).permute(2,0,1).unsqueeze(0).float() |
| if pixels.shape[-2:]!=(h*16,w*16): |
| pixels=torch.nn.functional.interpolate(pixels,size=(h*16,w*16),mode='bicubic',align_corners=False,antialias=True) |
| mean=torch.tensor([.48145466,.4578275,.40821073]).view(1,3,1,1) |
| std=torch.tensor([.26862954,.26130258,.27577711]).view(1,3,1,1) |
| return SimpleNamespace(pixel_values=(pixels/255.-mean)/std) |
|
|
