<jupyter_start><jupyter_text>Transform Definitions<jupyter_code>class UnNormalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, tensor):
        """
        Args:
            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
        Returns:
            Tensor: Normalized image.
        """
        for t, m, s in zip(tensor, self.mean, self.std):
            t.mul_(s).add_(m)
        return tensor

    
transform = T.Compose(
    [
        T.Resize((128, 128)),
        T.ToTensor(),
        T.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])
    ]
)<jupyter_output><empty_output><jupyter_text>Weight Paths and Image PathsWeight paths<jupyter_code>ROOT_LFW = Path("datasets/eval/data/lfw/lfw_masked")

WEIGHT_PATH = "weights/mask_exp19-resnetSTN/epoch_18/mask_exp19-resnetSTN_ep18.pth"

P1 = "Alfredo_Moreno/Alfredo_Moreno_0001.jpg"
P2 = "Dyab_Abou_Jahjah/Dyab_Abou_Jahjah_0001.jpg"
def stn_forward(weight_path, img_path):
    net = ResnetFaceSTN(stn_mode='resnet')
    net.load_state_dict(torch.load(weight_path))
    net.eval()

    unorm = UnNormalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])

    img = PIL.Image.open(img_path)
    img_t = transform(img).unsqueeze(0)
    img_stn = net.stn(img_t)
    img_p = kornia.tensor_to_image(unorm(img_stn.clone().detach()))
    
    return img_p<jupyter_output><empty_output><jupyter_text>Preview<jupyter_code>img_p = stn_forward(WEIGHT_PATH, ROOT_LFW / P1)
plt.imshow(img_p)
plt.show()
# plt.imsave("pair_sample_result/NEG_mCASIAC_P1_2.jpg", img_p)

img_p = stn_forward(WEIGHT_PATH, ROOT_LFW / P2)
plt.imshow(img_p)
plt.show()
# plt.imsave("pair_sample_result/NEG_mCASIAC_P2_2.jpg", img_p)
sample_images_masked = glob.glob('image_samples/LFW-Masked/*.jpg')
sample_images = glob.glob('image_samples/LFW/*.jpg')

top_row = []
bottom_row = []
for i, imgs in enumerate(zip(sample_images_masked, sample_images), 1):
    masked_img, norm_img = imgs
    
    # img_a = stn_forward(WEIGHT_PATH, masked_img)
    img_a = stn_forward(WEIGHT_PATH, norm_img)

    # img_o = plt.imread(masked_img)
    img_o = plt.imread(norm_img)
    img_o = resize(np.array(img_o), (112, 112), anti_aliasing=True)

    if len(top_row) == 0: top_row = img_o
    else: top_row = np.hstack((top_row, img_o))
    
    if len(bottom_row) == 0: bottom_row = np.array(img_a)
    else: bottom_row = np.hstack((bottom_row, img_a))
    
aligned_img = np.vstack((top_row, bottom_row))
plt.imsave("pair_sample_result/nonmask_ori_align.jpg", aligned_img)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13914728682170543, 'fraction_numerical': 0.024806201550387597, 'mean_word_length': 3.8790170132325144, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '666836', 'license': 'MIT', 'n_tokens_mistral': 1108, 'n_tokens_neox': 1067, 'n_words': 155, 'path': 'get_stn_alignment.ipynb', 'repo_name': 'agikarasugi/Face-Mask-Invariant-End-to-End-Face-Recognition', 'size': '2580'}