| 77 | return x if x.shape[-2:] == CLIP_IMAGE_SIZE else F.interpolate(x, size=CLIP_IMAGE_SIZE, mode='bilinear', align_corners=True, antialias=True) |
| 78 | |
| 79 | class MsgDataset(Dataset): |
| 80 | |
| 81 | def __init__(self, **params): |
| 82 | for k, v in params.items(): |
| 83 | self.__setattr__(k, v) |
| 84 | |
| 85 | self.load_data() |
| 86 | |
| 87 | def __getitem__(self, index): |
| 88 | return self.data[index] |
| 89 | |
| 90 | def __len__(self): |
| 91 | return len(self.data) |
| 92 | |
| 93 | def load_data(self): |
| 94 | max_value = int(2 ** self.msg_len) |
| 95 | if not hasattr(self, 'b2t_maps'): |
| 96 | rand_tokens = torch.randperm(CLIP_TOKEN_MAX) + 1 |
| 97 | self.b2t_maps = rand_tokens[:2 * self.msg_len].reshape(self.msg_len, 2) |
| 98 | |
| 99 | if not hasattr(self, 'data'): |
| 100 | dec_messages = sample_K_from_N(self.max_size, max_value) |
| 101 | self.data = [self.dec_message_to_tokens(x) for x in dec_messages] |
| 102 | |
| 103 | # Decimal message -> Binary message + CLIP tokens |
| 104 | def dec_message_to_tokens(self, dec_message): |
| 105 | # Decimal message -> Binary message |
| 106 | bin_message_text = bin(dec_message)[2:].zfill(self.msg_len) |
| 107 | bin_message = torch.tensor([int(x) for x in bin_message_text]) |
| 108 | |
| 109 | # Binary message -> Tokens |
| 110 | tokens = torch.zeros(CLIP_TOKEN_LEN) |
| 111 | tokens[0] = CLIP_TOKEN_BEGIN |
| 112 | tokens[self.msg_len + 1] = CLIP_TOKEN_END |
| 113 | for idx, (bit, b2t_map )in enumerate(zip(bin_message, self.b2t_maps)): |
| 114 | tokens[idx + 1] = b2t_map[bit] |
| 115 | |
| 116 | return tokens.long(), bin_message.float() |
| 117 | |
| 118 | def get_params(self): |
| 119 | return {k : getattr(self, k) for k in ['data', 'b2t_maps']} |
| 120 | |
| 121 | # sample K Decimal messages from N population |
| 122 | def sample_K_from_N(K, N, Limit=48): |