Spaces:
Runtime error
Runtime error
| diff --git a/dataset/caption_dataset.py b/dataset/caption_dataset.py | |
| index 266fdda..0cc5d3f 100644 | |
| --- a/dataset/caption_dataset.py | |
| +++ b/dataset/caption_dataset.py | |
| class Caption(Dataset): | |
| elif self.dataset == 'demo': | |
| img_path_split = self.data_list[index]['image'].split('/') | |
| img_name = img_path_split[-2] + '/' + img_path_split[-1] | |
| - image, labels, labels_info = get_expert_labels('', self.label_path, img_name, 'helpers', self.experts) | |
| + image, labels, labels_info = get_expert_labels('prismer', self.label_path, img_name, 'helpers', self.experts) | |
| experts = self.transform(image, labels) | |
| experts = post_label_process(experts, labels_info) | |
| diff --git a/dataset/utils.py b/dataset/utils.py | |
| index b368aac..418358c 100644 | |
| --- a/dataset/utils.py | |
| +++ b/dataset/utils.py | |
| # https://github.com/NVlabs/prismer/blob/main/LICENSE | |
| import os | |
| +import pathlib | |
| import re | |
| import json | |
| import torch | |
| import torchvision.transforms as transforms | |
| import torchvision.transforms.functional as transforms_f | |
| from dataset.randaugment import RandAugment | |
| -COCO_FEATURES = torch.load('dataset/coco_features.pt')['features'] | |
| -ADE_FEATURES = torch.load('dataset/ade_features.pt')['features'] | |
| -DETECTION_FEATURES = torch.load('dataset/detection_features.pt')['features'] | |
| -BACKGROUND_FEATURES = torch.load('dataset/background_features.pt') | |
| +cur_dir = pathlib.Path(__file__).parent | |
| + | |
| +COCO_FEATURES = torch.load(cur_dir / 'coco_features.pt')['features'] | |
| +ADE_FEATURES = torch.load(cur_dir / 'ade_features.pt')['features'] | |
| +DETECTION_FEATURES = torch.load(cur_dir / 'detection_features.pt')['features'] | |
| +BACKGROUND_FEATURES = torch.load(cur_dir / 'background_features.pt') | |
| class Transform: | |
| diff --git a/model/prismer.py b/model/prismer.py | |
| index 080253a..02362a4 100644 | |
| --- a/model/prismer.py | |
| +++ b/model/prismer.py | |
| # https://github.com/NVlabs/prismer/blob/main/LICENSE | |
| import json | |
| +import pathlib | |
| import torch.nn as nn | |
| from model.modules.vit import load_encoder | |
| from model.modules.roberta import load_decoder | |
| from transformers import RobertaTokenizer, RobertaConfig | |
| +cur_dir = pathlib.Path(__file__).parent | |
| + | |
| + | |
| class Prismer(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| class Prismer(nn.Module): | |
| elif exp in ['obj_detection', 'ocr_detection']: | |
| self.experts[exp] = 64 | |
| - prismer_config = json.load(open('configs/prismer.json', 'r'))[config['prismer_model']] | |
| + prismer_config = json.load(open(f'{cur_dir.parent}/configs/prismer.json', 'r'))[config['prismer_model']] | |
| roberta_config = RobertaConfig.from_dict(prismer_config['roberta_model']) | |
| self.tokenizer = RobertaTokenizer.from_pretrained(prismer_config['roberta_model']['model_name']) | |
| class Prismer(nn.Module): | |
| self.prepare_to_train(config['freeze']) | |
| self.ignored_modules = self.get_ignored_modules(config['freeze']) | |
| - | |
| + | |
| def prepare_to_train(self, mode='none'): | |
| for name, params in self.named_parameters(): | |
| if mode == 'freeze_lang': | |