| |
| |
| |
|
|
| from pdb import set_trace as bb |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import torchvision.transforms as tvf |
|
|
| from core.conv_mixer import ConvMixer |
|
|
| norm_RGB = tvf.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
|
|
|
|
| class PixelDesc (nn.Module): |
| def __init__(self, path='models/PUMP_st.pt'): |
| super().__init__() |
| state_dict = torch.load( path, 'cpu' ) |
| self.pixel_desc = ConvMixer(output_dim=128, hidden_dim=512, depth=7, patch_size=4, kernel_size=9).eval() |
| self.pixel_desc.load_state_dict(state_dict) |
|
|
| def configure(self, pipeline): |
| |
| pipeline.__class__ = type(type(pipeline).__name__+'_Trained', (DescPipeline, type(pipeline)), {}) |
| return self |
|
|
| def get_atomic_patch_size(self): |
| return 4 |
|
|
| def forward(self, img, stride=1, offset=0): |
| if img.ndim == 3: img = img[None] |
| trf = torch.eye(3, device=img.device) |
|
|
| desc = self.pixel_desc( img ) |
| desc = desc[..., offset::stride, offset::stride].contiguous() |
| return desc, trf |
|
|
|
|
| class DescPipeline: |
| def extract_descs(self, img1, img2, dtype=None): |
| |
| img1, sca1 = self.demultiplex_img_trf(img1) |
| img2, sca2 = self.demultiplex_img_trf(img2) |
|
|
| |
| fimg1, fimg2 = [norm_RGB(img.type(dtype)/255) for img in (img1, img2)] |
|
|
| self.pixel_desc.type(fimg1.dtype) |
| desc1, trf1 = self.pixel_desc(fimg1, stride=4, offset=2) |
| desc2, trf2 = self.pixel_desc(fimg2) |
| return (img1, img2), (desc1.type(dtype), desc2.type(dtype)), (sca1@trf1, sca2@trf2) |
|
|
| def first_level(self, desc1, desc2, **kw): |
| B, C, H, W = desc1.shape |
| weights = desc1.permute(0, 2, 3, 1).view(H*W, C, 1, 1) |
| corr = F.conv2d(desc2, weights, padding=0, bias=None)[0] |
| norms = torch.ones(desc1.shape[-2:], device=corr.device) |
| return corr.view(desc1.shape[-2:]+desc2.shape[-2:]), norms |
|
|
