test commit

.gitattributes

@@ -25,3 +25,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+
+*.pth filter=lfs diff=lfs merge=lfs -text

StyTR.py

@@ -0,0 +1,230 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+import numpy as np
+import box_ops
+from misc import (NestedTensor, nested_tensor_from_tensor_list,
+                       accuracy, get_world_size, interpolate,
+                       is_dist_avail_and_initialized)
+from function import normal,normal_style
+from function import calc_mean_std
+import scipy.stats as stats
+from ViT_helper import DropPath, to_2tuple, trunc_normal_
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=256, patch_size=8, in_chans=3, embed_dim=512):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+        
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        self.up1 = nn.Upsample(scale_factor=2, mode='nearest')
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x)
+
+        return x
+
+
+decoder = nn.Sequential(
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 256, (3, 3)),
+    nn.ReLU(),
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 128, (3, 3)),
+    nn.ReLU(),
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 128, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 64, (3, 3)),
+    nn.ReLU(),
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 64, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 3, (3, 3)),
+)
+
+vgg = nn.Sequential(
+    nn.Conv2d(3, 3, (1, 1)),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(3, 64, (3, 3)),
+    nn.ReLU(),  # relu1-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 64, (3, 3)),
+    nn.ReLU(),  # relu1-2
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 128, (3, 3)),
+    nn.ReLU(),  # relu2-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 128, (3, 3)),
+    nn.ReLU(),  # relu2-2
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 256, (3, 3)),
+    nn.ReLU(),  # relu3-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),  # relu3-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),  # relu3-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),  # relu3-4
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 512, (3, 3)),
+    nn.ReLU(),  # relu4-1, this is the last layer used
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-4
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU()  # relu5-4
+)
+
+class MLP(nn.Module):
+    """ Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+class StyTrans(nn.Module):
+    """ This is the style transform transformer module """
+    
+    def __init__(self,encoder,decoder,PatchEmbed, transformer,args):
+
+        super().__init__()
+        enc_layers = list(encoder.children())
+        self.enc_1 = nn.Sequential(*enc_layers[:4])  # input -> relu1_1
+        self.enc_2 = nn.Sequential(*enc_layers[4:11])  # relu1_1 -> relu2_1
+        self.enc_3 = nn.Sequential(*enc_layers[11:18])  # relu2_1 -> relu3_1
+        self.enc_4 = nn.Sequential(*enc_layers[18:31])  # relu3_1 -> relu4_1
+        self.enc_5 = nn.Sequential(*enc_layers[31:44])  # relu4_1 -> relu5_1
+        
+        for name in ['enc_1', 'enc_2', 'enc_3', 'enc_4', 'enc_5']:
+            for param in getattr(self, name).parameters():
+                param.requires_grad = False
+
+        self.mse_loss = nn.MSELoss()
+        self.transformer = transformer
+        hidden_dim = transformer.d_model       
+        self.decode = decoder
+        self.embedding = PatchEmbed
+
+    def encode_with_intermediate(self, input):
+        results = [input]
+        for i in range(5):
+            func = getattr(self, 'enc_{:d}'.format(i + 1))
+            results.append(func(results[-1]))
+        return results[1:]
+
+    def calc_content_loss(self, input, target):
+      assert (input.size() == target.size())
+      assert (target.requires_grad is False)
+      return self.mse_loss(input, target)
+
+    def calc_style_loss(self, input, target):
+        assert (input.size() == target.size())
+        assert (target.requires_grad is False)
+        input_mean, input_std = calc_mean_std(input)
+        target_mean, target_std = calc_mean_std(target)
+        return self.mse_loss(input_mean, target_mean) + \
+               self.mse_loss(input_std, target_std)
+    def forward(self, samples_c: NestedTensor,samples_s: NestedTensor):
+        """ The forward expects a NestedTensor, which consists of:
+               - samples.tensor: batched images, of shape [batch_size x 3 x H x W]
+               - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels
+
+        """
+        content_input = samples_c
+        style_input = samples_s
+        if isinstance(samples_c, (list, torch.Tensor)):
+            samples_c = nested_tensor_from_tensor_list(samples_c)   # support different-sized images padding is used for mask [tensor, mask] 
+        if isinstance(samples_s, (list, torch.Tensor)):
+            samples_s = nested_tensor_from_tensor_list(samples_s) 
+        
+        # ### features used to calcate loss 
+        content_feats = self.encode_with_intermediate(samples_c.tensors)
+        style_feats = self.encode_with_intermediate(samples_s.tensors)
+
+        ### Linear projection
+        style = self.embedding(samples_s.tensors)
+        content = self.embedding(samples_c.tensors)
+        
+        # postional embedding is calculated in transformer.py
+        pos_s = None
+        pos_c = None
+
+        mask = None
+        hs = self.transformer(style, mask , content, pos_c, pos_s)   
+        Ics = self.decode(hs)
+
+        Ics_feats = self.encode_with_intermediate(Ics)
+        loss_c = self.calc_content_loss(normal(Ics_feats[-1]), normal(content_feats[-1]))+self.calc_content_loss(normal(Ics_feats[-2]), normal(content_feats[-2]))
+        # Style loss
+        loss_s = self.calc_style_loss(Ics_feats[0], style_feats[0])
+        for i in range(1, 5):
+            loss_s += self.calc_style_loss(Ics_feats[i], style_feats[i])
+            
+        
+        Icc = self.decode(self.transformer(content, mask , content, pos_c, pos_c))
+        Iss = self.decode(self.transformer(style, mask , style, pos_s, pos_s))    
+
+        #Identity losses lambda 1    
+        loss_lambda1 = self.calc_content_loss(Icc,content_input)+self.calc_content_loss(Iss,style_input)
+        
+        #Identity losses lambda 2
+        Icc_feats=self.encode_with_intermediate(Icc)
+        Iss_feats=self.encode_with_intermediate(Iss)
+        loss_lambda2 = self.calc_content_loss(Icc_feats[0], content_feats[0])+self.calc_content_loss(Iss_feats[0], style_feats[0])
+        for i in range(1, 5):
+            loss_lambda2 += self.calc_content_loss(Icc_feats[i], content_feats[i])+self.calc_content_loss(Iss_feats[i], style_feats[i])
+        # Please select and comment out one of the following two sentences
+        return Ics,  loss_c, loss_s, loss_lambda1, loss_lambda2   #train
+        # return Ics    #test 

sofaApp.py

@@ -0,0 +1,50 @@
+import numpy as np
+import gradio as gr
+from Segmentation.segmentation import get_mask,replace_sofa
+from StyleTransfer.styleTransfer import resize_sofa,resize_style,create_styledSofa
+from PIL import Image
+
+def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
+    """
+    Styles (all) the sofas in the image to the given style.
+    This function uses a transformer to combine the image with the desired style according 
+    to a generated mask of the sofas in the image.
+    Input:
+        input_img = image containing a sofa
+        style_img = image containing a style
+    Return:
+        new_sofa  = image containing the styled sofa
+    """
+    
+    # preprocess input images to be (640,640) squares to fit requirements of the segmentation model
+    resized_img = resize_sofa(input_img)
+    resized_style = resize_style(style_img)
+    # generate mask for image
+    mask = get_mask(resized_img)
+    styled_sofa = create_styledSofa(resized_img,resized_style)
+    new_sofa = replace_sofa(resized_img,mask,styled_sofa)
+    return new_sofa
+
+image = gr.inputs.Image()
+style = gr.inputs.Image()
+
+demo = gr.Interface(
+    style_sofa,
+    [image,style],
+    'image',
+    examples=[
+        ['input/sofa_example1.jpg','input/style_example1.jpg'],
+        ['input/sofa_example1.jpg','input/style_example2.jpg'],
+        ['input/sofa_example1.jpg','input/style_example3.jpg'],
+        ['input/sofa_example1.jpg','input/style_example4.jpg'],
+        ['input/sofa_example1.jpg','input/style_example5.jpg'],
+    ],
+    title="Style your sofa",
+    description="🛋 Customize your sofa to your wildest dreams! 🛋",
+)
+
+if __name__ == "__main__":
+    demo.launch(share=True)
+
+
+#https://github.com/dhawan98/Post-Processing-of-Image-Segmentation-using-CRF

test.py

@@ -0,0 +1,176 @@
+import argparse
+from pathlib import Path
+import os
+import torch
+import torch.nn as nn
+from PIL import Image
+from os.path import basename
+from os.path import splitext
+from torchvision import transforms
+from torchvision.utils import save_image
+from function import calc_mean_std, normal, coral
+import transformer as transformer
+import StyTR as StyTR
+import matplotlib.pyplot as plt
+from matplotlib import cm
+from function import normal
+import numpy as np
+
+def test_transform(size, crop):
+    transform_list = []
+   
+    if size != 0: 
+        transform_list.append(transforms.Resize(size))
+    if crop:
+        transform_list.append(transforms.CenterCrop(size))
+    transform_list.append(transforms.ToTensor())
+    transform = transforms.Compose(transform_list)
+    return transform
+def style_transform(h,w):
+    k = (h,w)
+    size = int(np.max(k))
+    transform_list = []    
+    transform_list.append(transforms.CenterCrop((h,w)))
+    transform_list.append(transforms.ToTensor())
+    transform = transforms.Compose(transform_list)
+    return transform
+
+def content_transform():
+    
+    transform_list = []   
+    transform_list.append(transforms.ToTensor())
+    transform = transforms.Compose(transform_list)
+    return transform
+
+  
+
+parser = argparse.ArgumentParser()
+# Basic options
+parser.add_argument('--content', type=str,
+                    help='File path to the content image')
+parser.add_argument('--content_dir', type=str,
+                    help='Directory path to a batch of content images')
+parser.add_argument('--style', type=str,
+                    help='File path to the style image, or multiple style \
+                    images separated by commas if you want to do style \
+                    interpolation or spatial control')
+parser.add_argument('--style_dir', type=str,
+                    help='Directory path to a batch of style images')
+parser.add_argument('--output', type=str, default='output',
+                    help='Directory to save the output image(s)')
+parser.add_argument('--vgg', type=str, default='./experiments/vgg_normalised.pth')
+parser.add_argument('--decoder_path', type=str, default='experiments/decoder_iter_160000.pth')
+parser.add_argument('--Trans_path', type=str, default='experiments/transformer_iter_160000.pth')
+parser.add_argument('--embedding_path', type=str, default='experiments/embedding_iter_160000.pth')
+
+
+parser.add_argument('--style_interpolation_weights', type=str, default="")
+parser.add_argument('--a', type=float, default=1.0)
+parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'),
+                        help="Type of positional embedding to use on top of the image features")
+parser.add_argument('--hidden_dim', default=512, type=int,
+                        help="Size of the embeddings (dimension of the transformer)")
+args = parser.parse_args()
+
+# Advanced options
+content_size=640
+style_size=640
+crop='store_true'
+save_ext='.jpg'
+output_path=args.output
+preserve_color='store_true'
+alpha=args.a
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Either --content or --content_dir should be given.
+if args.content:
+    content_paths = [Path(args.content)]
+else:
+    content_dir = Path(args.content_dir)
+    content_paths = [f for f in content_dir.glob('*')]
+
+# Either --style or --style_dir should be given.
+if args.style:
+    style_paths = [Path(args.style)]    
+else:
+    style_dir = Path(args.style_dir)
+    style_paths = [f for f in style_dir.glob('*')]
+
+if not os.path.exists(output_path):
+    os.mkdir(output_path)
+
+
+vgg = StyTR.vgg
+vgg.load_state_dict(torch.load(args.vgg))
+vgg = nn.Sequential(*list(vgg.children())[:44])
+
+decoder = StyTR.decoder
+Trans = transformer.Transformer()
+embedding = StyTR.PatchEmbed()
+
+decoder.eval()
+Trans.eval()
+vgg.eval()
+from collections import OrderedDict
+new_state_dict = OrderedDict()
+state_dict = torch.load(args.decoder_path)
+for k, v in state_dict.items():
+    #namekey = k[7:] # remove `module.`
+    namekey = k
+    new_state_dict[namekey] = v
+decoder.load_state_dict(new_state_dict)
+
+new_state_dict = OrderedDict()
+state_dict = torch.load(args.Trans_path)
+for k, v in state_dict.items():
+    #namekey = k[7:] # remove `module.`
+    namekey = k
+    new_state_dict[namekey] = v
+Trans.load_state_dict(new_state_dict)
+
+new_state_dict = OrderedDict()
+state_dict = torch.load(args.embedding_path)
+for k, v in state_dict.items():
+    #namekey = k[7:] # remove `module.`
+    namekey = k
+    new_state_dict[namekey] = v
+embedding.load_state_dict(new_state_dict)
+
+network = StyTR.StyTrans(vgg,decoder,embedding,Trans,args)
+network.eval()
+network.to(device)
+
+
+
+content_tf = test_transform(content_size, crop)
+style_tf = test_transform(style_size, crop)
+
+for content_path in content_paths:
+    for style_path in style_paths:
+        print(content_path)
+       
+      
+        content_tf1 = content_transform()       
+        content = content_tf(Image.open(content_path).convert("RGB"))
+
+        h,w,c=np.shape(content)    
+        style_tf1 = style_transform(h,w)
+        style = style_tf(Image.open(style_path).convert("RGB"))
+
+      
+        style = style.to(device).unsqueeze(0)
+        content = content.to(device).unsqueeze(0)
+        
+        with torch.no_grad():
+            output= network(content,style)     
+        output = output[0].cpu()
+                
+        output_name = '{:s}/{:s}_stylized_{:s}{:s}'.format(
+            output_path, splitext(basename(content_path))[0],
+            splitext(basename(style_path))[0], save_ext
+        )
+ 
+        save_image(output, output_name)
+   
+

transformer.py

@@ -0,0 +1,322 @@
+import copy
+from typing import Optional, List
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+from function import normal,normal_style
+import numpy as np
+import os
+device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
+os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"
+class Transformer(nn.Module):
+
+    def __init__(self, d_model=512, nhead=8, num_encoder_layers=3,
+                 num_decoder_layers=3, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False,
+                 return_intermediate_dec=False):
+        super().__init__()
+
+        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
+                                                dropout, activation, normalize_before)
+        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
+        self.encoder_c = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+        self.encoder_s = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+
+        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
+                                                dropout, activation, normalize_before)
+        decoder_norm = nn.LayerNorm(d_model)
+        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
+                                          return_intermediate=return_intermediate_dec)
+
+        self._reset_parameters()
+
+        self.d_model = d_model
+        self.nhead = nhead
+
+        self.new_ps = nn.Conv2d(512 , 512 , (1,1))
+        self.averagepooling = nn.AdaptiveAvgPool2d(18)
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, style, mask , content, pos_embed_c, pos_embed_s):
+
+        # content-aware positional embedding
+        content_pool = self.averagepooling(content)       
+        pos_c = self.new_ps(content_pool)
+        pos_embed_c = F.interpolate(pos_c, mode='bilinear',size= style.shape[-2:])
+
+        ###flatten NxCxHxW to HWxNxC     
+        style = style.flatten(2).permute(2, 0, 1)
+        if pos_embed_s is not None:
+            pos_embed_s = pos_embed_s.flatten(2).permute(2, 0, 1)
+      
+        content = content.flatten(2).permute(2, 0, 1)
+        if pos_embed_c is not None:
+            pos_embed_c = pos_embed_c.flatten(2).permute(2, 0, 1)
+     
+        
+        style = self.encoder_s(style, src_key_padding_mask=mask, pos=pos_embed_s)
+        content = self.encoder_c(content, src_key_padding_mask=mask, pos=pos_embed_c)
+        hs = self.decoder(content, style, memory_key_padding_mask=mask,
+                          pos=pos_embed_s, query_pos=pos_embed_c)[0]
+        
+        ### HWxNxC to NxCxHxW to
+        N, B, C= hs.shape          
+        H = int(np.sqrt(N))
+        hs = hs.permute(1, 2, 0)
+        hs = hs.view(B, C, -1,H)
+
+        return hs
+
+
+class TransformerEncoder(nn.Module):
+
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super().__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(self, src,
+                mask: Optional[Tensor] = None,
+                src_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None):
+        output = src
+        
+        for layer in self.layers:
+            output = layer(output, src_mask=mask,
+                           src_key_padding_mask=src_key_padding_mask, pos=pos)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+
+class TransformerDecoder(nn.Module):
+
+    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        output = tgt
+
+        intermediate = []
+
+        for layer in self.layers:
+            output = layer(output, memory, tgt_mask=tgt_mask,
+                           memory_mask=memory_mask,
+                           tgt_key_padding_mask=tgt_key_padding_mask,
+                           memory_key_padding_mask=memory_key_padding_mask,
+                           pos=pos, query_pos=query_pos)
+            if self.return_intermediate:
+                intermediate.append(self.norm(output))
+
+        if self.norm is not None:
+            output = self.norm(output)
+            if self.return_intermediate:
+                intermediate.pop()
+                intermediate.append(output)
+
+        if self.return_intermediate:
+            return torch.stack(intermediate)
+
+        return output.unsqueeze(0)
+
+
+class TransformerEncoderLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self,
+                     src,
+                     src_mask: Optional[Tensor] = None,
+                     src_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None):
+        q = k = self.with_pos_embed(src, pos)
+        # q = k = src
+        # print(q.size(),k.size(),src.size())
+        src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        return src
+
+    def forward_pre(self, src,
+                    src_mask: Optional[Tensor] = None,
+                    src_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None):
+        src2 = self.norm1(src)
+        q = k = self.with_pos_embed(src2, pos)
+        src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src2 = self.norm2(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
+        src = src + self.dropout2(src2)
+        return src
+
+    def forward(self, src,
+                src_mask: Optional[Tensor] = None,
+                src_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
+        return self.forward_post(src, src_mask, src_key_padding_mask, pos)
+
+
+class TransformerDecoderLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False):
+        super().__init__()
+        # d_model embedding dim
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self, tgt, memory,
+                     tgt_mask: Optional[Tensor] = None,
+                     memory_mask: Optional[Tensor] = None,
+                     tgt_key_padding_mask: Optional[Tensor] = None,
+                     memory_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None,
+                     query_pos: Optional[Tensor] = None):
+
+       
+        q = self.with_pos_embed(tgt, query_pos)
+        k = self.with_pos_embed(memory, pos)
+        v = memory 
+ 
+        tgt2 = self.self_attn(q, k, v, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+    
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
+                                   key=self.with_pos_embed(memory, pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+    def forward_pre(self, tgt, memory,
+                    tgt_mask: Optional[Tensor] = None,
+                    memory_mask: Optional[Tensor] = None,
+                    tgt_key_padding_mask: Optional[Tensor] = None,
+                    memory_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None,
+                    query_pos: Optional[Tensor] = None):
+        tgt2 = self.norm1(tgt)
+        q = k = self.with_pos_embed(tgt2, query_pos)
+        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+
+        tgt = tgt + self.dropout1(tgt2)
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
+                                   key=self.with_pos_embed(memory, pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
+                                    tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
+        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
+                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def build_transformer(args):
+    return Transformer(
+        d_model=args.hidden_dim,
+        dropout=args.dropout,
+        nhead=args.nheads,
+        dim_feedforward=args.dim_feedforward,
+        num_encoder_layers=args.enc_layers,
+        num_decoder_layers=args.dec_layers,
+        normalize_before=args.pre_norm,
+        return_intermediate_dec=True,
+    )
+
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")