File size: 9,476 Bytes
3b175de
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
import json


def get_prompt_list(file_path: str) -> list[list[str]]:

        
    return json.load(open(file_path, "r"))

class GetPromptList(object):
    _SUPPORTED_SOURCE = {'Sachit-descriptors',}
    def __init__(self, file_path: str, name2idx: dict[str: int] = None, class_names: list[str] = None) -> None:
        self.class_names = class_names
        self.file_path = file_path
        self.desc = get_prompt_list(file_path)
        if isinstance(self.desc, dict):
            self.__get_parts()
        if name2idx is not None:
            self.name2idx = name2idx
        elif class_names is not None:
            self.name2idx = {cls_name: idx for idx, cls_name in enumerate(class_names)}
        else:
            self.name2idx = {cls_name: idx for idx, cls_name in enumerate(self.desc.keys())} if isinstance(self.desc, dict) else None
        
    
    # def __get_sachit_desc(self, file_path: str): 
    #     params = get_sachit_hparams(file_path)
    #     return load_gpt_descriptions(params)
    
    def __get_parts(self, ):
        # get part names from one of the descriptions
        self.part_names = [d.split(":")[0].strip() for d in self.desc[list(self.desc.keys())[0]]]
    
    
    @staticmethod
    def replace_class_names(self, descs: dict, target_class: list[str], new_classes: list[str]):
        new_descs = []
        for desc, cls_name, new_name in zip(descs, target_class, new_classes):
            temp = [d.replace(cls_name, new_name) for d in desc]
            new_descs.extend(temp)
        return new_descs
    
    def __call__(self, source: str, pad: bool = False, max_len: int = 15, pad_text: str = "", target_classes: list[int] = None, pad_neg_index: bool = True):
        """
        This function will return a list of prompts based on the source (format) and file_path provided. 
        If name2idx is provided, the prompts will be mapped based on the provied class indexes. Otherwise,
        the prompts will be mapped based on the order of class name in the file.
        Note: this function is will apply trucation when padding is True to make sure to have fixed length prompts.

        Args:
            source (str): The sorce (format) of the prompts. Supported sources are: {self._SUPPORTED_SOURCE}
            file_path (str): The file that contains the original prompts.
            pad (bool, optional): Whether to pad the prompts to the same length. Defaults to False.
            max_len (int, optional): The maximum length of the prompts. Defaults to 15.
            pad_text (str, optional): The text to pad the prompts. Defaults to "Padding".
            target_classes (list[int], optional): A list of class indexes to include in the prompts. Defaults to None (include all classes).
        Returns:
            prompts (list[str]): A list of engineered prompts.
            class_idxs (list[int]): A list of class indexes for each prompt.
            class_mapping (dict[int: str]): A mapping of class indexes to class names.
        """
        org_desc_mapper = None
        match source:
            case 'Sachit-descriptors':
                desc, org_dict = self.__get_sachit_desc(self.file_path)

            case 'Sachit-no-template':
                desc = self.desc
                
            case 'Sachit-CLIP-template-5':
                desc, org_dict = self.__get_sachit_desc(self.file_path)
                desc = {k: [f'a photo of a {d}.' for d in v] for k, v in desc.items()}
            case 'cub-12-parts':
                return self.desc, None, None, None
            case 'chatgpt-no-template':
                desc = self.desc
            case 'chatgpt-template-0':
                # convert 'part: features' to 'a features part'
                template = 'a {} {}.'
                desc = {k: [template.format(d.split(":")[1].strip(), d.split(":")[0].strip()) for d in v] for k, v in self.desc.items()}
            case 'chatgpt-template-8':
                # convert '{part}: {features}' to 'a {features} {part} of {class_name}'
                template = 'a {} {} of {}.'
                desc = {k: [template.format(d.split(":")[1].strip(), d.split(":")[0].strip(), k) for d in v] for k, v in self.desc.items()}
            case 'chatgpt-template-5':
                # convert '{part}: {features}' to 'a photo of {class_name}, which is/has/etc {descriptor}
                desc, org_dict = self.__get_sachit_desc(self.file_path)
                template = 'a photo of a {}'
                desc = {k: [template.format(d) for d in v] for k, v in desc.items()}
            case 'chatgpt-template-x':
                # convert '{part}: {features}' to '{features}. {part}. {class_name}'
                desc = {k: [f'{d.split(":")[1].strip()}. {d.split(":")[0].strip()}. {k}' for d in v] for k, v in self.desc.items()}
            case 'chatgpt-template-x-2': # no class name
                # convert '{part}: {features}' to 'a {features} {part}'
                desc = {k: [f'a {d.split(":")[1].strip()} {d.split(":")[0].strip()}' for d in v] for k, v in self.desc.items()}
            case 'chatgpt-template-x-3': # no class name
                # convert '{part}: {features}' to '{features}. {part}.'
                desc = {k: [f'{d.split(":")[1].strip()}. {d.split(":")[0].strip()}.' for d in v] for k, v in self.desc.items()}
            case 'chatgpt-template-x-4': 
                # convert '{part}: {features}' to 'a {part} of {class_name}: {features}'
                desc = {k: [f'a {d.split(":")[0].strip()} of {k}: {d.split(":")[1].strip()}' for d in v] for k, v in self.desc.items()}
            
            case _:
                raise ValueError(f"Source {source} is not supported. Check {self._SUPPORTED_SOURCE}")
        
        # get the subset of descriptrions that match the target classes
        if len(self.name2idx) < len(desc):
            desc = {k: desc[k] for k in self.name2idx}     
            
        prompts, class_idxs, class_list = [], [], []
        class_mapping = {v: k for k, v in self.name2idx.items()}
        for class_name, class_idx in self.name2idx.items():
            descriptions = desc[class_name]
            if target_classes is not None and class_idx not in target_classes:
                continue
            if pad:
                pad_id = -1 if pad_neg_index else class_idx
                ids = [class_idx] * len(descriptions) + [pad_id] * (max_len - len(descriptions)) if len(descriptions) < max_len else [class_idx] * max_len
                if len(descriptions) < max_len:
                    descriptions.extend([pad_text] * (max_len - len(descriptions)))
                else:
                    descriptions = descriptions[:max_len]
            else:
                ids = [class_idx] * len(descriptions)
            prompts.extend(descriptions)
            class_idxs.extend(ids)
            class_list.append(class_name)
        
            if org_desc_mapper is not None:
                org_desc_mapper = {des: org_dict[class_name][des] for des in descriptions}

                
        return prompts, class_idxs, class_mapping, org_desc_mapper, class_list
    
    
imagenet_templates = [
    'a bad photo of a {}.',
    'a photo of many {}.',
    'a sculpture of a {}.',
    'a photo of the hard to see {}.',
    'a low resolution photo of the {}.',
    'a rendering of a {}.',
    'graffiti of a {}.',
    'a bad photo of the {}.',
    'a cropped photo of the {}.',
    'a tattoo of a {}.',
    'the embroidered {}.',
    'a photo of a hard to see {}.',
    'a bright photo of a {}.',
    'a photo of a clean {}.',
    'a photo of a dirty {}.',
    'a dark photo of the {}.',
    'a drawing of a {}.',
    'a photo of my {}.',
    'the plastic {}.',
    'a photo of the cool {}.',
    'a close-up photo of a {}.',
    'a black and white photo of the {}.',
    'a painting of the {}.',
    'a painting of a {}.',
    'a pixelated photo of the {}.',
    'a sculpture of the {}.',
    'a bright photo of the {}.',
    'a cropped photo of a {}.',
    'a plastic {}.',
    'a photo of the dirty {}.',
    'a jpeg corrupted photo of a {}.',
    'a blurry photo of the {}.',
    'a photo of the {}.',
    'a good photo of the {}.',
    'a rendering of the {}.',
    'a {} in a video game.',
    'a photo of one {}.',
    'a doodle of a {}.',
    'a close-up photo of the {}.',
    'a photo of a {}.',
    'the origami {}.',
    'the {} in a video game.',
    'a sketch of a {}.',
    'a doodle of the {}.',
    'a origami {}.',
    'a low resolution photo of a {}.',
    'the toy {}.',
    'a rendition of the {}.',
    'a photo of the clean {}.',
    'a photo of a large {}.',
    'a rendition of a {}.',
    'a photo of a nice {}.',
    'a photo of a weird {}.',
    'a blurry photo of a {}.',
    'a cartoon {}.',
    'art of a {}.',
    'a sketch of the {}.',
    'a embroidered {}.',
    'a pixelated photo of a {}.',
    'itap of the {}.',
    'a jpeg corrupted photo of the {}.',
    'a good photo of a {}.',
    'a plushie {}.',
    'a photo of the nice {}.',
    'a photo of the small {}.',
    'a photo of the weird {}.',
    'the cartoon {}.',
    'art of the {}.',
    'a drawing of the {}.',
    'a photo of the large {}.',
    'a black and white photo of a {}.',
    'the plushie {}.',
    'a dark photo of a {}.',
    'itap of a {}.',
    'graffiti of the {}.',
    'a toy {}.',
    'itap of my {}.',
    'a photo of a cool {}.',
    'a photo of a small {}.',
    'a tattoo of the {}.',
]