Add all files

app.py

@@ -0,0 +1,120 @@
+from urllib.request import urlopen
+import argparse
+import clip
+from PIL import Image
+import pandas as pd
+import time
+import torch
+from dataloader.extract_features_dataloader import transform_resize, question_preprocess
+from model.vqa_model import NetVQA
+from dataclasses import dataclass
+from torch.cuda.amp import autocast
+import gradio as gr
+
+@dataclass
+class InferenceConfig:
+    '''
+    Describes configuration of the training process
+    '''
+    model: str = "RN50x64"
+    checkpoint_root_clip: str = "./checkpoints/clip"
+    checkpoint_root_head: str = "./checkpoints/head"
+    
+    use_question_preprocess: bool = True   # True: delete ? at end
+    
+    aux_mapping = {0: "unanswerable",
+                   1: "unsuitable",
+                   2: "yes",
+                   3: "no",
+                   4: "number",
+                   5: "color",
+                   6: "other"}
+    folds = 10   
+    tta = False
+    # Data
+    n_classes: int =  5726
+    
+    # class mapping
+    class_mapping: str = "./data/annotations/class_mapping.csv"
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+         
+config = InferenceConfig()
+
+# load class mapping
+cm = pd.read_csv(config.class_mapping)
+classid_to_answer = {}
+for i in range(len(cm)):
+    row = cm.iloc[i]
+    classid_to_answer[row["class_id"]] = row["answer"]
+
+clip_model, preprocess = clip.load(config.model, download_root=config.checkpoint_root_clip)
+
+model = NetVQA(config).to(config.device)
+
+
+config.checkpoint_head = "{}/{}.pt".format(config.checkpoint_root_head, config.model)
+
+model_state_dict = torch.load(config.checkpoint_head)  
+model.load_state_dict(model_state_dict, strict=True)
+
+
+#%%
+# Select Preprocessing
+image_transforms = transform_resize(clip_model.visual.input_resolution)
+
+if config.use_question_preprocess:
+    question_transforms = question_preprocess
+else:
+    question_transforms = None
+
+clip_model.eval()
+
+
+def predict(img, text):
+    img = Image.fromarray(img)
+    if config.tta:   
+        image_augmentations = []
+        for transform in image_transforms:
+            image_augmentations.append(transform(img))
+        img = torch.stack(image_augmentations, dim=0)
+    else:
+        img = image_transforms(img)
+        img = img.unsqueeze(dim=0)
+
+    question = question_transforms(text)
+    question_tokens = clip.tokenize(question, truncate=True)
+    with torch.no_grad():
+        img = img.to(config.device)
+        img_feature = clip_model.encode_image(img)
+        if config.tta:
+            weights = torch.tensor(config.features_selection).reshape((len(config.features_selection),1))    
+            img_feature =  img_feature *  weights.to(config.device)
+            img_feature = img_feature.sum(0)
+            img_feature = img_feature.unsqueeze(0)
+        
+        question_tokens = question_tokens.to(config.device)
+        question_feature = clip_model.encode_text(question_tokens)
+
+        with autocast():
+            output, output_aux = model(img_feature, question_feature)
+
+    prediction_vqa = dict()
+    output = output.cpu().squeeze(0)
+    for k, v in classid_to_answer.items():
+        prediction_vqa[v] = float(output[k])
+
+    prediction_aux = dict()
+    output_aux = output_aux.cpu().squeeze(0)
+    for k, v in config.aux_mapping.items():
+        prediction_aux[v] = float(output_aux[k])
+
+
+    return prediction_vqa, prediction_aux
+
+gr.Interface(fn=predict, 
+             inputs=[gr.Image(label='Image'), gr.Textbox(label='Question')],
+             outputs=[gr.outputs.Label(label='Answer', num_top_classes=5), gr.outputs.Label(label='Answer Category', num_top_classes=7)],
+             examples=[['examples/VizWiz_train_00004056.jpg', 'Is that a beer or a coke?'], ['examples/VizWiz_train_00017146.jpg', 'Can you tell me what\'s on this envelope please?'], ['examples/VizWiz_val_00003077.jpg', 'What is this?']]
+             ).launch()
+

dataloader/extract_features_dataloader.py

@@ -0,0 +1,268 @@
+import pandas as pd
+import os
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+import clip
+from torch.utils.data import DataLoader
+import torchvision.transforms as tf
+import torchvision.transforms.functional as TF
+
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+    
+    
+class ExtractFeaturesDataset(Dataset):
+    def __init__(self,
+                 annotations,
+                 img_path,
+                 image_transforms=None,
+                 question_transforms=None,
+                 tta=False):
+        
+        
+        self.img_path = img_path
+        self.image_transforms = image_transforms
+        self.question_transforms = question_transforms
+        
+        self.img_ids = annotations["image_id"].values
+        self.split = annotations["split"].values
+        self.questions = annotations["question"].values
+        
+        self.tta = tta
+       
+
+
+    def __getitem__(self, index):
+
+        image_id = self.img_ids[index]
+        split = self.split[index]
+        
+        # image input
+        with open(os.path.join(self.img_path, split, image_id), "rb") as f:
+            img = Image.open(f)
+            
+            if self.tta:   
+                image_augmentations = []
+                
+                for transform in self.image_transforms:
+                    
+                    image_augmentations.append(transform(img))
+               
+        
+                img = torch.stack(image_augmentations, dim=0)
+                
+            else:
+                img = self.image_transforms(img)
+            
+        question = self.questions[index]
+        
+        if self.question_transforms: 
+            question = self.question_transforms(question)
+
+        # question input
+        question = clip.tokenize(question, truncate=True) 
+        question = question.squeeze()
+            
+        return img, question, image_id
+    
+    def __len__(self):
+        return len(self.img_ids)
+    
+        
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def Sharpen(sharpness_factor=1.0):
+    
+    def wrapper(x):
+    
+        return TF.adjust_sharpness(x, sharpness_factor)
+    
+    return wrapper
+
+
+def Rotate(angle=0.0):
+    
+    def wrapper(x):
+        return TF.rotate(x, angle)
+    
+    return wrapper
+
+def transform_crop(n_px):
+    return tf.Compose([
+        tf.Resize(n_px, interpolation=BICUBIC),
+        tf.CenterCrop(n_px),
+        _convert_image_to_rgb,
+        tf.ToTensor(),
+        tf.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+def transform_crop_rotate(n_px, rotation_angle=0.0):
+    return tf.Compose([
+        Rotate(angle=rotation_angle),
+        tf.Resize(n_px, interpolation=BICUBIC),
+        tf.CenterCrop(n_px),
+        _convert_image_to_rgb,
+        tf.ToTensor(),
+        tf.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+def transform_resize(n_px):
+    return tf.Compose([
+        tf.Resize((n_px, n_px), interpolation=BICUBIC),
+        _convert_image_to_rgb,
+        tf.ToTensor(),
+        tf.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+def transform_resize_rotate(n_px, rotation_angle=0.0):
+    return tf.Compose([
+        Rotate(angle=rotation_angle),
+        tf.Resize((n_px, n_px), interpolation=BICUBIC),
+        _convert_image_to_rgb,
+        tf.ToTensor(),
+        tf.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+def get_tta_preprocess(img_size):
+    
+    img_preprocess = [
+                      transform_crop(img_size),
+                      transform_crop_rotate(img_size, rotation_angle=90.0),
+                      transform_crop_rotate(img_size, rotation_angle=270.0),
+                      transform_resize(img_size),
+                      transform_resize_rotate(img_size, rotation_angle=90.0),
+                      transform_resize_rotate(img_size, rotation_angle=270.0),
+                      ]
+    
+    return img_preprocess
+
+def question_preprocess(question, debug=False):
+    
+    question = question.replace("?", ".")
+
+    if question[-1] == " ":
+        question = question[:-1]
+         
+
+    if question[-1] != ".":
+        question = question + "."				
+				   
+    if debug:
+        print("Question:", question)
+        
+    return question
+
+            
+def get_dataloader_extraction(config):
+
+
+    if config.use_question_preprocess:
+        print("Using custom preprocessing: Question")
+        question_transforms = question_preprocess
+    else:
+        question_transforms = None
+    
+    if config.tta:
+        ("Using augmentation transforms:")
+        img_preprocess = get_tta_preprocess(config.img_size)
+    else:
+        ("Using original CLIP transforms:")
+        img_preprocess = transform_crop(config.img_size)
+           
+    
+
+    train_data = pd.read_csv(config.train_annotations_path)
+      
+    train_dataset = ExtractFeaturesDataset(annotations = train_data,
+                                      img_path=config.img_path,
+                                      image_transforms=img_preprocess,
+                                      question_transforms=question_transforms,
+                                      tta=config.tta)
+    
+    
+    
+    train_loader = DataLoader(dataset=train_dataset, 
+                              batch_size=config.batch_size, 
+                              shuffle=False,
+                              num_workers=config.num_workers)
+    
+         
+    
+    test_data = pd.read_csv(config.test_annotations_path)
+      
+    test_dataset = ExtractFeaturesDataset(annotations = test_data,
+                                      img_path=config.img_path,
+                                      image_transforms=img_preprocess,
+                                      question_transforms=question_transforms,
+                                      tta=config.tta)
+    
+    
+    test_loader = ExtractFeaturesDataset(dataset=test_dataset, 
+                              batch_size=config.batch_size, 
+                              shuffle=False,
+                              num_workers=config.num_workers)
+    
+    return train_loader, test_loader
+
+
+def get_dataloader_inference(config):
+    
+    if config.use_question_preprocess:
+        print("Using custom preprocessing: Question")
+        question_transforms = question_preprocess
+    else:
+        question_transforms = None
+    
+    if config.tta:
+        ("Using augmentation transforms:")
+        img_preprocess = transform_resize(config.img_size)
+    else:
+        ("Using original CLIP transforms:")
+        img_preprocess = transform_crop(config.img_size)
+           
+    
+
+    train_data = pd.read_csv(config.train_annotations_path)
+      
+    train_dataset = ExtractFeaturesDataset(annotations = train_data,
+                                      img_path=config.img_path,
+                                      image_transforms=img_preprocess,
+                                      question_transforms=question_transforms,
+                                      tta=config.tta)
+    
+    
+    
+    train_loader = DataLoader(dataset=train_dataset, 
+                              batch_size=config.batch_size, 
+                              shuffle=False,
+                              num_workers=config.num_workers)
+    
+         
+    
+    test_data = pd.read_csv(config.test_annotations_path)
+      
+    test_dataset = ExtractFeaturesDataset(annotations = test_data,
+                                      img_path=config.img_path,
+                                      image_transforms=img_preprocess,
+                                      question_transforms=question_transforms,
+                                      tta=config.tta)
+    
+    
+    test_loader = ExtractFeaturesDataset(dataset=test_dataset, 
+                              batch_size=config.batch_size, 
+                              shuffle=False,
+                              num_workers=config.num_workers)
+    
+    return train_loader, test_loader
+    
+         
+    

model/vqa_model.py

@@ -0,0 +1,123 @@
+import torch
+
+class HeadVQA(torch.nn.Module):
+    def __init__(self, train_config):
+        super().__init__()
+        
+        embedding_size = {'RN50': 1024,
+                          'RN101': 512,
+                          'RN50x4': 640,
+                          'RN50x16': 768,
+                          'RN50x64': 1024,
+                          'ViT-B/32': 512,
+                          'ViT-B/16': 512,
+                          'ViT-L/14': 768,
+                          'ViT-L/14@336px': 768}
+        
+        n_aux_classes = len(set(train_config.aux_mapping.values()))
+        
+        self.ln1 = torch.nn.LayerNorm(embedding_size[train_config.model]*2)
+        self.dp1 = torch.nn.Dropout(0.5)
+        self.fc1 = torch.nn.Linear(embedding_size[train_config.model] * 2, 512)
+
+        self.ln2 = torch.nn.LayerNorm(512)
+        self.dp2 = torch.nn.Dropout(0.5)
+        self.fc2 = torch.nn.Linear(512, train_config.n_classes)
+        
+        self.fc_aux = torch.nn.Linear(512, n_aux_classes)
+        self.fc_gate = torch.nn.Linear(n_aux_classes, train_config.n_classes)
+        self.act_gate = torch.nn.Sigmoid()
+        
+
+    def forward(self, img_features, question_features):
+        xc = torch.cat((img_features, question_features), dim=-1) 
+        
+        x = self.ln1(xc)
+        x = self.dp1(x)
+        x = self.fc1(x)
+        
+        aux = self.fc_aux(x)
+
+        gate = self.fc_gate(aux)
+        gate = self.act_gate(gate)
+         
+        x = self.ln2(x)
+        x = self.dp2(x)
+        vqa = self.fc2(x)
+        
+        output = vqa * gate
+        
+        return output, aux
+    
+    
+class NetVQA(torch.nn.Module):
+    def __init__(self, train_config):
+        super().__init__()
+        
+        self.heads = torch.nn.ModuleList()
+        
+        if isinstance(train_config.folds, list):
+            self.num_heads = len(train_config.folds)
+        else:
+            self.num_heads = train_config.folds
+        
+        for i in range(self.num_heads):
+            self.heads.append(HeadVQA(train_config))
+        
+        
+    def forward(self, img_features, question_features):
+        
+        output = []
+        output_aux = []
+        
+        for head in self.heads:
+            
+            logits, logits_aux = head(img_features, question_features)
+            
+            probs = logits.softmax(-1)
+            probs_aux = logits_aux.softmax(-1)
+            
+            output.append(probs)
+            output_aux.append(probs_aux)
+             
+        output = torch.stack(output, dim=-1).mean(-1)
+        output_aux = torch.stack(output_aux, dim=-1).mean(-1)
+                
+        return output, output_aux  
+      
+def merge_vqa(train_config):
+    
+    # Initialize model
+    model = NetVQA(train_config)
+
+
+    for fold in train_config.folds:
+        
+        print("load weights from fold {} into head {}".format(fold, fold))
+        
+        checkpoint_path = "{}/{}/fold_{}".format(train_config.model_path, train_config.model, fold)
+        
+        if train_config.crossvalidation:
+            # load best checkpoint 
+            model_state_dict = torch.load('{}/weights_best.pth'.format(checkpoint_path))      
+        else:
+            # load checkpoint on train end
+            model_state_dict = torch.load('{}/weights_end.pth'.format(checkpoint_path))    
+             
+        model.heads[fold].load_state_dict(model_state_dict, strict=True)
+        
+    checkpoint_path = "{}/{}/weights_merged.pth".format(train_config.model_path, train_config.model)
+
+    print("Saving weights of merged model:", checkpoint_path)
+
+    torch.save(model.state_dict(), checkpoint_path) 
+    
+    return model
+
+
+
+        
+        
+        
+        
+