(self, img)
| 63 | self.to_pixels = nn.Linear(decoder_dim, pixel_values_per_patch) |
| 64 | |
| 65 | def forward(self, img): |
| 66 | |
| 67 | # get patches and their number |
| 68 | patches = self.to_patch(img) |
| 69 | print('patch shape:', patches.shape) |
| 70 | _, num_patches, *_ = patches.shape |
| 71 | |
| 72 | # project pixel patches to tokens and add positions |
| 73 | tokens = self.patch_to_emb(patches) |
| 74 | tokens = tokens + self.encoder.pos_embedding[:, 1:(num_patches + 1)] |
| 75 | |
| 76 | # encode tokens by the encoder |
| 77 | encoded_tokens = self.encoder.transformer(tokens) |
| 78 | |
| 79 | # project encoder to decoder dimensions, if they are not equal. |
| 80 | decoder_tokens = self.enc_to_dec(encoded_tokens) |
| 81 | |
| 82 | # decode tokens with decoder |
| 83 | decoded_tokens = self.decoder(decoder_tokens) |
| 84 | |
| 85 | # project tokens to pixels |
| 86 | pred_pixel_values = self.to_pixels(decoded_tokens) |
| 87 | |
| 88 | # reshape |
| 89 | # being consistent with the input parameters of function "build_model" |
| 90 | # p1 = patch_size, p2 = patch_size, h=image_size[0]//patch_size) |
| 91 | rec_images = rearrange(pred_pixel_values, 'b (h w) (p1 p2 c) -> b c (h p1) (w p2)', |
| 92 | p1 = 16, p2 = 16, h=256//16) |
| 93 | |
| 94 | return rec_images, patches |
| 95 | |
| 96 | def build_model(setting , image_size, patch_size): |
| 97 | """ |
nothing calls this directly
no outgoing calls
no test coverage detected