Spaces:
Running
on
A10G
Running
on
A10G
| import torch | |
| from lib.utils.iimage import IImage | |
| class InputMask: | |
| def to(self, device): return InputMask(self.image, device = device) | |
| def cuda(self): return InputMask(self.image, device = 'cuda') | |
| def cpu(self): return InputMask(self.image, device = 'cpu') | |
| def __init__(self, input_image, device = 'cpu'): | |
| ''' | |
| args: | |
| input_image: (b,c,h,w) tensor | |
| ''' | |
| if hasattr(input_image, 'is_iimage'): | |
| self.image = input_image | |
| self.val512 = self.full = (input_image.torch(0) > 0.5).float() | |
| elif isinstance(input_image, torch.Tensor): | |
| self.val512 = self.full = input_image.clone() | |
| self.image = IImage(input_image,0) | |
| self.h,self.w = h,w = self.val512.shape[-2:] | |
| self.shape = [self.h, self.w] | |
| self.shape64 = [self.h // 8, self.w // 8] | |
| self.shape32 = [self.h // 16, self.w // 16] | |
| self.shape16 = [self.h // 32, self.w // 32] | |
| self.shape8 = [self.h // 64, self.w // 64] | |
| self.res = self.h * self.w | |
| self.res64 = self.res // 64 | |
| self.res32 = self.res // 64 // 4 | |
| self.res16 = self.res // 64 // 16 | |
| self.res8 = self.res // 64 // 64 | |
| self.img = self.image | |
| self.img512 = self.image | |
| self.img64 = self.image.resize((h//8,w//8)) | |
| self.img32 = self.image.resize((h//16,w//16)) | |
| self.img16 = self.image.resize((h//32,w//32)) | |
| self.img8 = self.image.resize((h//64,w//64)) | |
| self.val64 = (self.img64.torch(0) > 0.5).float() | |
| self.val32 = (self.img32.torch(0) > 0.5).float() | |
| self.val16 = (self.img16.torch(0) > 0.5).float() | |
| self.val8 = ( self.img8.torch(0) > 0.5).float() | |
| def get_res(self, q, device = 'cpu'): | |
| if q.shape[1] == self.res64: return self.val64.to(device) | |
| if q.shape[1] == self.res32: return self.val32.to(device) | |
| if q.shape[1] == self.res16: return self.val16.to(device) | |
| if q.shape[1] == self.res8: return self.val8.to(device) | |
| def get_res(self, q, device = 'cpu'): | |
| if q.shape[1] == self.res64: return self.val64.to(device) | |
| if q.shape[1] == self.res32: return self.val32.to(device) | |
| if q.shape[1] == self.res16: return self.val16.to(device) | |
| if q.shape[1] == self.res8: return self.val8.to(device) | |
| def get_shape(self, q, device = 'cpu'): | |
| if q.shape[1] == self.res64: return self.shape64 | |
| if q.shape[1] == self.res32: return self.shape32 | |
| if q.shape[1] == self.res16: return self.shape16 | |
| if q.shape[1] == self.res8: return self.shape8 | |
| def get_res_val(self, q, device = 'cpu'): | |
| if q.shape[1] == self.res64: return 64 | |
| if q.shape[1] == self.res32: return 32 | |
| if q.shape[1] == self.res16: return 16 | |
| if q.shape[1] == self.res8: return 8 | |
| class InputMask2: | |
| def to(self, device): return InputMask2(self.image, device = device) | |
| def cuda(self): return InputMask2(self.image, device = 'cuda') | |
| def cpu(self): return InputMask2(self.image, device = 'cpu') | |
| def __init__(self, input_image, device = 'cpu'): | |
| ''' | |
| args: | |
| input_image: (b,c,h,w) tensor | |
| ''' | |
| if hasattr(input_image, 'is_iimage'): | |
| self.image = input_image | |
| self.val512 = self.full = input_image.torch(0).bool().float() | |
| elif isinstance(input_image, torch.Tensor): | |
| self.val512 = self.full = input_image.clone() | |
| self.image = IImage(input_image,0) | |
| self.h,self.w = h,w = self.val512.shape[-2:] | |
| self.shape = [self.h, self.w] | |
| self.shape64 = [self.h // 8, self.w // 8] | |
| self.shape32 = [self.h // 16, self.w // 16] | |
| self.shape16 = [self.h // 32, self.w // 32] | |
| self.shape8 = [self.h // 64, self.w // 64] | |
| self.res = self.h * self.w | |
| self.res64 = self.res // 64 | |
| self.res32 = self.res // 64 // 4 | |
| self.res16 = self.res // 64 // 16 | |
| self.res8 = self.res // 64 // 64 | |
| self.img = self.image | |
| self.img512 = self.image | |
| self.img64 = self.image.resize((h//8,w//8)) | |
| self.img32 = self.image.resize((h//16,w//16)) | |
| self.img16 = self.image.resize((h//32,w//32)).dilate(1) | |
| self.img8 = self.image.resize((h//64,w//64)).dilate(1) | |
| self.val64 = self.img64.torch(0).bool().float() | |
| self.val32 = self.img32.torch(0).bool().float() | |
| self.val16 = self.img16.torch(0).bool().float() | |
| self.val8 = self.img8.torch(0).bool().float() | |
| def get_res(self, q, device = 'cpu'): | |
| if q.shape[1] == self.res64: return self.val64.to(device) | |
| if q.shape[1] == self.res32: return self.val32.to(device) | |
| if q.shape[1] == self.res16: return self.val16.to(device) | |
| if q.shape[1] == self.res8: return self.val8.to(device) | |
| def get_res(self, q, device = 'cpu'): | |
| if q.shape[1] == self.res64: return self.val64.to(device) | |
| if q.shape[1] == self.res32: return self.val32.to(device) | |
| if q.shape[1] == self.res16: return self.val16.to(device) | |
| if q.shape[1] == self.res8: return self.val8.to(device) | |
| def get_shape(self, q, device = 'cpu'): | |
| if q.shape[1] == self.res64: return self.shape64 | |
| if q.shape[1] == self.res32: return self.shape32 | |
| if q.shape[1] == self.res16: return self.shape16 | |
| if q.shape[1] == self.res8: return self.shape8 | |
| def get_res_val(self, q, device = 'cpu'): | |
| if q.shape[1] == self.res64: return 64 | |
| if q.shape[1] == self.res32: return 32 | |
| if q.shape[1] == self.res16: return 16 | |
| if q.shape[1] == self.res8: return 8 |