wufan commited on
Commit
18e4b60
1 Parent(s): c0d1798

Upload 111 files

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +1 -0
  2. examples/0000004.png +0 -0
  3. examples/0000005.png +0 -0
  4. examples/0000006.png +0 -0
  5. examples/0000007.png +0 -0
  6. examples/0000011.png +0 -0
  7. unimernet/__init__.py +31 -0
  8. unimernet/__pycache__/__init__.cpython-310.pyc +0 -0
  9. unimernet/common/__init__.py +0 -0
  10. unimernet/common/__pycache__/__init__.cpython-310.pyc +0 -0
  11. unimernet/common/__pycache__/config.cpython-310.pyc +0 -0
  12. unimernet/common/__pycache__/dist_utils.cpython-310.pyc +0 -0
  13. unimernet/common/__pycache__/logger.cpython-310.pyc +0 -0
  14. unimernet/common/__pycache__/registry.cpython-310.pyc +0 -0
  15. unimernet/common/__pycache__/utils.cpython-310.pyc +0 -0
  16. unimernet/common/config.py +468 -0
  17. unimernet/common/dist_utils.py +137 -0
  18. unimernet/common/gradcam.py +24 -0
  19. unimernet/common/logger.py +195 -0
  20. unimernet/common/optims.py +120 -0
  21. unimernet/common/registry.py +329 -0
  22. unimernet/common/utils.py +424 -0
  23. unimernet/configs/datasets/formula/formula_eval.yaml +6 -0
  24. unimernet/configs/datasets/formula/formula_train.yaml +6 -0
  25. unimernet/configs/datasets/formula/multi_scale_formula_train.yaml +21 -0
  26. unimernet/configs/default.yaml +10 -0
  27. unimernet/configs/models/unimernet_base.yaml +31 -0
  28. unimernet/datasets/__init__.py +0 -0
  29. unimernet/datasets/__pycache__/__init__.cpython-310.pyc +0 -0
  30. unimernet/datasets/__pycache__/data_utils.cpython-310.pyc +0 -0
  31. unimernet/datasets/builders/__init__.py +69 -0
  32. unimernet/datasets/builders/__pycache__/__init__.cpython-310.pyc +0 -0
  33. unimernet/datasets/builders/__pycache__/base_dataset_builder.cpython-310.pyc +0 -0
  34. unimernet/datasets/builders/__pycache__/formula.cpython-310.pyc +0 -0
  35. unimernet/datasets/builders/base_dataset_builder.py +233 -0
  36. unimernet/datasets/builders/formula.py +105 -0
  37. unimernet/datasets/data_utils.py +284 -0
  38. unimernet/datasets/datasets/__pycache__/base_dataset.cpython-310.pyc +0 -0
  39. unimernet/datasets/datasets/__pycache__/formula.cpython-310.pyc +0 -0
  40. unimernet/datasets/datasets/__pycache__/formula_multi_scale.cpython-310.pyc +0 -0
  41. unimernet/datasets/datasets/base_dataset.py +103 -0
  42. unimernet/datasets/datasets/dataloader_utils.py +200 -0
  43. unimernet/datasets/datasets/formula.py +71 -0
  44. unimernet/datasets/datasets/formula_multi_scale.py +32 -0
  45. unimernet/models/__init__.py +198 -0
  46. unimernet/models/__pycache__/__init__.cpython-310.pyc +0 -0
  47. unimernet/models/__pycache__/base_model.cpython-310.pyc +0 -0
  48. unimernet/models/__pycache__/clip_vit.cpython-310.pyc +0 -0
  49. unimernet/models/__pycache__/eva_vit.cpython-310.pyc +0 -0
  50. unimernet/models/base_model.py +251 -0
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ unimernet/processors/formula_processor_helper/frost/frost1.png filter=lfs diff=lfs merge=lfs -text
examples/0000004.png ADDED
examples/0000005.png ADDED
examples/0000006.png ADDED
examples/0000007.png ADDED
examples/0000011.png ADDED
unimernet/__init__.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ import os
9
+ import sys
10
+
11
+ from omegaconf import OmegaConf
12
+
13
+ from unimernet.common.registry import registry
14
+
15
+ from unimernet.datasets.builders import *
16
+ from unimernet.models import *
17
+ from unimernet.processors import *
18
+ from unimernet.tasks import *
19
+
20
+
21
+ root_dir = os.path.dirname(os.path.abspath(__file__))
22
+ default_cfg = OmegaConf.load(os.path.join(root_dir, "configs/default.yaml"))
23
+
24
+ registry.register_path("library_root", root_dir)
25
+ repo_root = os.path.join(root_dir, "..")
26
+ registry.register_path("repo_root", repo_root)
27
+ cache_root = os.path.join(repo_root, default_cfg.env.cache_root)
28
+ registry.register_path("cache_root", cache_root)
29
+
30
+ registry.register("MAX_INT", sys.maxsize)
31
+ registry.register("SPLIT_NAMES", ["train", "val", "test"])
unimernet/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (1.02 kB). View file
 
unimernet/common/__init__.py ADDED
File without changes
unimernet/common/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (160 Bytes). View file
 
unimernet/common/__pycache__/config.cpython-310.pyc ADDED
Binary file (12.1 kB). View file
 
unimernet/common/__pycache__/dist_utils.cpython-310.pyc ADDED
Binary file (3.77 kB). View file
 
unimernet/common/__pycache__/logger.cpython-310.pyc ADDED
Binary file (6.42 kB). View file
 
unimernet/common/__pycache__/registry.cpython-310.pyc ADDED
Binary file (8.39 kB). View file
 
unimernet/common/__pycache__/utils.cpython-310.pyc ADDED
Binary file (12.8 kB). View file
 
unimernet/common/config.py ADDED
@@ -0,0 +1,468 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ import logging
9
+ import json
10
+ from typing import Dict
11
+
12
+ from omegaconf import OmegaConf
13
+ from unimernet.common.registry import registry
14
+
15
+
16
+ class Config:
17
+ def __init__(self, args):
18
+ self.config = {}
19
+
20
+ self.args = args
21
+
22
+ # Register the config and configuration for setup
23
+ registry.register("configuration", self)
24
+
25
+ user_config = self._build_opt_list(self.args.options)
26
+
27
+ config = OmegaConf.load(self.args.cfg_path)
28
+
29
+ runner_config = self.build_runner_config(config)
30
+ model_config = self.build_model_config(config, **user_config)
31
+ dataset_config = self.build_dataset_config(config)
32
+
33
+ # Validate the user-provided runner configuration
34
+ # model and dataset configuration are supposed to be validated by the respective classes
35
+ # [TODO] validate the model/dataset configuration
36
+ # self._validate_runner_config(runner_config)
37
+
38
+ # Override the default configuration with user options.
39
+ self.config = OmegaConf.merge(
40
+ runner_config, model_config, dataset_config, user_config
41
+ )
42
+
43
+ def _validate_runner_config(self, runner_config):
44
+ """
45
+ This method validates the configuration, such that
46
+ 1) all the user specified options are valid;
47
+ 2) no type mismatches between the user specified options and the config.
48
+ """
49
+ runner_config_validator = create_runner_config_validator()
50
+ runner_config_validator.validate(runner_config)
51
+
52
+ def _build_opt_list(self, opts):
53
+ opts_dot_list = self._convert_to_dot_list(opts)
54
+ return OmegaConf.from_dotlist(opts_dot_list)
55
+
56
+ @staticmethod
57
+ def build_model_config(config, **kwargs):
58
+ model = config.get("model", None)
59
+ assert model is not None, "Missing model configuration file."
60
+
61
+ model_cls = registry.get_model_class(model.arch)
62
+ assert model_cls is not None, f"Model '{model.arch}' has not been registered."
63
+
64
+ model_type = kwargs.get("model.model_type", None)
65
+ if not model_type:
66
+ model_type = model.get("model_type", None)
67
+ # else use the model type selected by user.
68
+
69
+ assert model_type is not None, "Missing model_type."
70
+
71
+ model_config_path = model_cls.default_config_path(model_type=model_type)
72
+
73
+ model_config = OmegaConf.create()
74
+ # hiararchy override, customized config > default config
75
+ model_config = OmegaConf.merge(
76
+ model_config,
77
+ OmegaConf.load(model_config_path),
78
+ {"model": config["model"]},
79
+ )
80
+
81
+ return model_config
82
+
83
+ @staticmethod
84
+ def build_runner_config(config):
85
+ return {"run": config.run}
86
+
87
+ @staticmethod
88
+ def build_dataset_config(config):
89
+ datasets = config.get("datasets", None)
90
+ if datasets is None:
91
+ raise KeyError(
92
+ "Expecting 'datasets' as the root key for dataset configuration."
93
+ )
94
+
95
+ dataset_config = OmegaConf.create()
96
+
97
+ for dataset_name in datasets:
98
+ builder_cls = registry.get_builder_class(dataset_name)
99
+
100
+ dataset_config_type = datasets[dataset_name].get("type", "default")
101
+ dataset_config_path = builder_cls.default_config_path(
102
+ type=dataset_config_type
103
+ )
104
+
105
+ # hiararchy override, customized config > default config
106
+ dataset_config = OmegaConf.merge(
107
+ dataset_config,
108
+ OmegaConf.load(dataset_config_path),
109
+ {"datasets": {dataset_name: config["datasets"][dataset_name]}},
110
+ )
111
+
112
+ return dataset_config
113
+
114
+ def _convert_to_dot_list(self, opts):
115
+ if opts is None:
116
+ opts = []
117
+
118
+ if len(opts) == 0:
119
+ return opts
120
+
121
+ has_equal = opts[0].find("=") != -1
122
+
123
+ if has_equal:
124
+ return opts
125
+
126
+ return [(opt + "=" + value) for opt, value in zip(opts[0::2], opts[1::2])]
127
+
128
+ def get_config(self):
129
+ return self.config
130
+
131
+ @property
132
+ def run_cfg(self):
133
+ return self.config.run
134
+
135
+ @property
136
+ def datasets_cfg(self):
137
+ return self.config.datasets
138
+
139
+ @property
140
+ def model_cfg(self):
141
+ return self.config.model
142
+
143
+ def pretty_print(self):
144
+ logging.info("\n===== Running Parameters =====")
145
+ logging.info(self._convert_node_to_json(self.config.run))
146
+
147
+ logging.info("\n====== Dataset Attributes ======")
148
+ datasets = self.config.datasets
149
+
150
+ for dataset in datasets:
151
+ if dataset in self.config.datasets:
152
+ logging.info(f"\n======== {dataset} =======")
153
+ dataset_config = self.config.datasets[dataset]
154
+ logging.info(self._convert_node_to_json(dataset_config))
155
+ else:
156
+ logging.warning(f"No dataset named '{dataset}' in config. Skipping")
157
+
158
+ logging.info(f"\n====== Model Attributes ======")
159
+ logging.info(self._convert_node_to_json(self.config.model))
160
+
161
+ def _convert_node_to_json(self, node):
162
+ container = OmegaConf.to_container(node, resolve=True)
163
+ return json.dumps(container, indent=4, sort_keys=True)
164
+
165
+ def to_dict(self):
166
+ return OmegaConf.to_container(self.config)
167
+
168
+
169
+ def node_to_dict(node):
170
+ return OmegaConf.to_container(node)
171
+
172
+
173
+ class ConfigValidator:
174
+ """
175
+ This is a preliminary implementation to centralize and validate the configuration.
176
+ May be altered in the future.
177
+
178
+ A helper class to validate configurations from yaml file.
179
+
180
+ This serves the following purposes:
181
+ 1. Ensure all the options in the yaml are defined, raise error if not.
182
+ 2. when type mismatches are found, the validator will raise an error.
183
+ 3. a central place to store and display helpful messages for supported configurations.
184
+
185
+ """
186
+
187
+ class _Argument:
188
+ def __init__(self, name, choices=None, type=None, help=None):
189
+ self.name = name
190
+ self.val = None
191
+ self.choices = choices
192
+ self.type = type
193
+ self.help = help
194
+
195
+ def __str__(self):
196
+ s = f"{self.name}={self.val}"
197
+ if self.type is not None:
198
+ s += f", ({self.type})"
199
+ if self.choices is not None:
200
+ s += f", choices: {self.choices}"
201
+ if self.help is not None:
202
+ s += f", ({self.help})"
203
+ return s
204
+
205
+ def __init__(self, description):
206
+ self.description = description
207
+
208
+ self.arguments = dict()
209
+
210
+ self.parsed_args = None
211
+
212
+ def __getitem__(self, key):
213
+ assert self.parsed_args is not None, "No arguments parsed yet."
214
+
215
+ return self.parsed_args[key]
216
+
217
+ def __str__(self) -> str:
218
+ return self.format_help()
219
+
220
+ def add_argument(self, *args, **kwargs):
221
+ """
222
+ Assume the first argument is the name of the argument.
223
+ """
224
+ self.arguments[args[0]] = self._Argument(*args, **kwargs)
225
+
226
+ def validate(self, config=None):
227
+ """
228
+ Convert yaml config (dict-like) to list, required by argparse.
229
+ """
230
+ for k, v in config.items():
231
+ assert (
232
+ k in self.arguments
233
+ ), f"""{k} is not a valid argument. Support arguments are {self.format_arguments()}."""
234
+
235
+ if self.arguments[k].type is not None:
236
+ try:
237
+ self.arguments[k].val = self.arguments[k].type(v)
238
+ except ValueError:
239
+ raise ValueError(f"{k} is not a valid {self.arguments[k].type}.")
240
+
241
+ if self.arguments[k].choices is not None:
242
+ assert (
243
+ v in self.arguments[k].choices
244
+ ), f"""{k} must be one of {self.arguments[k].choices}."""
245
+
246
+ return config
247
+
248
+ def format_arguments(self):
249
+ return str([f"{k}" for k in sorted(self.arguments.keys())])
250
+
251
+ def format_help(self):
252
+ # description + key-value pair string for each argument
253
+ help_msg = str(self.description)
254
+ return help_msg + ", available arguments: " + self.format_arguments()
255
+
256
+ def print_help(self):
257
+ # display help message
258
+ print(self.format_help())
259
+
260
+
261
+ def create_runner_config_validator():
262
+ validator = ConfigValidator(description="Runner configurations")
263
+
264
+ validator.add_argument(
265
+ "runner",
266
+ type=str,
267
+ choices=["runner_base", "runner_iter"],
268
+ help="""Runner to use. The "runner_base" uses epoch-based training while iter-based
269
+ runner runs based on iters. Default: runner_base""",
270
+ )
271
+ # add argumetns for training dataset ratios
272
+ validator.add_argument(
273
+ "train_dataset_ratios",
274
+ type=Dict[str, float],
275
+ help="""Ratios of training dataset. This is used in iteration-based runner.
276
+ Do not support for epoch-based runner because how to define an epoch becomes tricky.
277
+ Default: None""",
278
+ )
279
+ validator.add_argument(
280
+ "max_iters",
281
+ type=float,
282
+ help="Maximum number of iterations to run.",
283
+ )
284
+ validator.add_argument(
285
+ "max_epoch",
286
+ type=int,
287
+ help="Maximum number of epochs to run.",
288
+ )
289
+ # add arguments for iters_per_inner_epoch
290
+ validator.add_argument(
291
+ "iters_per_inner_epoch",
292
+ type=float,
293
+ help="Number of iterations per inner epoch. This is required when runner is runner_iter.",
294
+ )
295
+ lr_scheds_choices = registry.list_lr_schedulers()
296
+ validator.add_argument(
297
+ "lr_sched",
298
+ type=str,
299
+ choices=lr_scheds_choices,
300
+ help="Learning rate scheduler to use, from {}".format(lr_scheds_choices),
301
+ )
302
+ task_choices = registry.list_tasks()
303
+ validator.add_argument(
304
+ "task",
305
+ type=str,
306
+ choices=task_choices,
307
+ help="Task to use, from {}".format(task_choices),
308
+ )
309
+ # add arguments for init_lr
310
+ validator.add_argument(
311
+ "init_lr",
312
+ type=float,
313
+ help="Initial learning rate. This will be the learning rate after warmup and before decay.",
314
+ )
315
+ # add arguments for min_lr
316
+ validator.add_argument(
317
+ "min_lr",
318
+ type=float,
319
+ help="Minimum learning rate (after decay).",
320
+ )
321
+ # add arguments for warmup_lr
322
+ validator.add_argument(
323
+ "warmup_lr",
324
+ type=float,
325
+ help="Starting learning rate for warmup.",
326
+ )
327
+ # add arguments for learning rate decay rate
328
+ validator.add_argument(
329
+ "lr_decay_rate",
330
+ type=float,
331
+ help="Learning rate decay rate. Required if using a decaying learning rate scheduler.",
332
+ )
333
+ # add arguments for weight decay
334
+ validator.add_argument(
335
+ "weight_decay",
336
+ type=float,
337
+ help="Weight decay rate.",
338
+ )
339
+ # add arguments for training batch size
340
+ validator.add_argument(
341
+ "batch_size_train",
342
+ type=int,
343
+ help="Training batch size.",
344
+ )
345
+ # add arguments for evaluation batch size
346
+ validator.add_argument(
347
+ "batch_size_eval",
348
+ type=int,
349
+ help="Evaluation batch size, including validation and testing.",
350
+ )
351
+ # add arguments for number of workers for data loading
352
+ validator.add_argument(
353
+ "num_workers",
354
+ help="Number of workers for data loading.",
355
+ )
356
+ # add arguments for warm up steps
357
+ validator.add_argument(
358
+ "warmup_steps",
359
+ type=int,
360
+ help="Number of warmup steps. Required if a warmup schedule is used.",
361
+ )
362
+ # add arguments for random seed
363
+ validator.add_argument(
364
+ "seed",
365
+ type=int,
366
+ help="Random seed.",
367
+ )
368
+ # add arguments for output directory
369
+ validator.add_argument(
370
+ "output_dir",
371
+ type=str,
372
+ help="Output directory to save checkpoints and logs.",
373
+ )
374
+ # add arguments for whether only use evaluation
375
+ validator.add_argument(
376
+ "evaluate",
377
+ help="Whether to only evaluate the model. If true, training will not be performed.",
378
+ )
379
+ # add arguments for splits used for training, e.g. ["train", "val"]
380
+ validator.add_argument(
381
+ "train_splits",
382
+ type=list,
383
+ help="Splits to use for training.",
384
+ )
385
+ # add arguments for splits used for validation, e.g. ["val"]
386
+ validator.add_argument(
387
+ "valid_splits",
388
+ type=list,
389
+ help="Splits to use for validation. If not provided, will skip the validation.",
390
+ )
391
+ # add arguments for splits used for testing, e.g. ["test"]
392
+ validator.add_argument(
393
+ "test_splits",
394
+ type=list,
395
+ help="Splits to use for testing. If not provided, will skip the testing.",
396
+ )
397
+ # add arguments for accumulating gradient for iterations
398
+ validator.add_argument(
399
+ "accum_grad_iters",
400
+ type=int,
401
+ help="Number of iterations to accumulate gradient for.",
402
+ )
403
+
404
+ # ====== distributed training ======
405
+ validator.add_argument(
406
+ "device",
407
+ type=str,
408
+ choices=["cpu", "cuda"],
409
+ help="Device to use. Support 'cuda' or 'cpu' as for now.",
410
+ )
411
+ validator.add_argument(
412
+ "world_size",
413
+ type=int,
414
+ help="Number of processes participating in the job.",
415
+ )
416
+ validator.add_argument("dist_url", type=str)
417
+ validator.add_argument("distributed", type=bool)
418
+ # add arguments to opt using distributed sampler during evaluation or not
419
+ validator.add_argument(
420
+ "use_dist_eval_sampler",
421
+ type=bool,
422
+ help="Whether to use distributed sampler during evaluation or not.",
423
+ )
424
+
425
+ # ====== task specific ======
426
+ # generation task specific arguments
427
+ # add arguments for maximal length of text output
428
+ validator.add_argument(
429
+ "max_len",
430
+ type=int,
431
+ help="Maximal length of text output.",
432
+ )
433
+ # add arguments for minimal length of text output
434
+ validator.add_argument(
435
+ "min_len",
436
+ type=int,
437
+ help="Minimal length of text output.",
438
+ )
439
+ # add arguments number of beams
440
+ validator.add_argument(
441
+ "num_beams",
442
+ type=int,
443
+ help="Number of beams used for beam search.",
444
+ )
445
+
446
+ # vqa task specific arguments
447
+ # add arguments for number of answer candidates
448
+ validator.add_argument(
449
+ "num_ans_candidates",
450
+ type=int,
451
+ help="""For ALBEF and BLIP, these models first rank answers according to likelihood to select answer candidates.""",
452
+ )
453
+ # add arguments for inference method
454
+ validator.add_argument(
455
+ "inference_method",
456
+ type=str,
457
+ choices=["genearte", "rank"],
458
+ help="""Inference method to use for question answering. If rank, requires a answer list.""",
459
+ )
460
+
461
+ # ====== model specific ======
462
+ validator.add_argument(
463
+ "k_test",
464
+ type=int,
465
+ help="Number of top k most similar samples from ITC/VTC selection to be tested.",
466
+ )
467
+
468
+ return validator
unimernet/common/dist_utils.py ADDED
@@ -0,0 +1,137 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ import datetime
9
+ import functools
10
+ import os
11
+
12
+ import torch
13
+ import torch.distributed as dist
14
+ import timm.models.hub as timm_hub
15
+
16
+
17
+ def setup_for_distributed(is_master):
18
+ """
19
+ This function disables printing when not in master process
20
+ """
21
+ import builtins as __builtin__
22
+
23
+ builtin_print = __builtin__.print
24
+
25
+ def print(*args, **kwargs):
26
+ force = kwargs.pop("force", False)
27
+ if is_master or force:
28
+ builtin_print(*args, **kwargs)
29
+
30
+ __builtin__.print = print
31
+
32
+
33
+ def is_dist_avail_and_initialized():
34
+ if not dist.is_available():
35
+ return False
36
+ if not dist.is_initialized():
37
+ return False
38
+ return True
39
+
40
+
41
+ def get_world_size():
42
+ if not is_dist_avail_and_initialized():
43
+ return 1
44
+ return dist.get_world_size()
45
+
46
+
47
+ def get_rank():
48
+ if not is_dist_avail_and_initialized():
49
+ return 0
50
+ return dist.get_rank()
51
+
52
+
53
+ def is_main_process():
54
+ return get_rank() == 0
55
+
56
+
57
+ def init_distributed_mode(args):
58
+ if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
59
+ args.rank = int(os.environ["RANK"])
60
+ args.world_size = int(os.environ["WORLD_SIZE"])
61
+ args.gpu = int(os.environ["LOCAL_RANK"])
62
+ elif "SLURM_PROCID" in os.environ:
63
+ args.rank = int(os.environ["SLURM_PROCID"])
64
+ args.gpu = args.rank % torch.cuda.device_count()
65
+ else:
66
+ print("Not using distributed mode")
67
+ args.distributed = False
68
+ return
69
+
70
+ args.distributed = True
71
+
72
+ torch.cuda.set_device(args.gpu)
73
+ args.dist_backend = "nccl"
74
+ print(
75
+ "| distributed init (rank {}, world {}): {}".format(
76
+ args.rank, args.world_size, args.dist_url
77
+ ),
78
+ flush=True,
79
+ )
80
+ torch.distributed.init_process_group(
81
+ backend=args.dist_backend,
82
+ init_method=args.dist_url,
83
+ world_size=args.world_size,
84
+ rank=args.rank,
85
+ timeout=datetime.timedelta(
86
+ days=365
87
+ ), # allow auto-downloading and de-compressing
88
+ )
89
+ torch.distributed.barrier()
90
+ setup_for_distributed(args.rank == 0)
91
+
92
+
93
+ def get_dist_info():
94
+ if torch.__version__ < "1.0":
95
+ initialized = dist._initialized
96
+ else:
97
+ initialized = dist.is_initialized()
98
+ if initialized:
99
+ rank = dist.get_rank()
100
+ world_size = dist.get_world_size()
101
+ else: # non-distributed training
102
+ rank = 0
103
+ world_size = 1
104
+ return rank, world_size
105
+
106
+
107
+ def main_process(func):
108
+ @functools.wraps(func)
109
+ def wrapper(*args, **kwargs):
110
+ rank, _ = get_dist_info()
111
+ if rank == 0:
112
+ return func(*args, **kwargs)
113
+
114
+ return wrapper
115
+
116
+
117
+ def download_cached_file(url, check_hash=True, progress=False):
118
+ """
119
+ Download a file from a URL and cache it locally. If the file already exists, it is not downloaded again.
120
+ If distributed, only the main process downloads the file, and the other processes wait for the file to be downloaded.
121
+ """
122
+
123
+ def get_cached_file_path():
124
+ # a hack to sync the file path across processes
125
+ parts = torch.hub.urlparse(url)
126
+ filename = os.path.basename(parts.path)
127
+ cached_file = os.path.join(timm_hub.get_cache_dir(), filename)
128
+
129
+ return cached_file
130
+
131
+ if is_main_process():
132
+ timm_hub.download_cached_file(url, check_hash, progress)
133
+
134
+ if is_dist_avail_and_initialized():
135
+ dist.barrier()
136
+
137
+ return get_cached_file_path()
unimernet/common/gradcam.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ from matplotlib import pyplot as plt
3
+ from scipy.ndimage import filters
4
+ from skimage import transform as skimage_transform
5
+
6
+
7
+ def getAttMap(img, attMap, blur=True, overlap=True):
8
+ attMap -= attMap.min()
9
+ if attMap.max() > 0:
10
+ attMap /= attMap.max()
11
+ attMap = skimage_transform.resize(attMap, (img.shape[:2]), order=3, mode="constant")
12
+ if blur:
13
+ attMap = filters.gaussian_filter(attMap, 0.02 * max(img.shape[:2]))
14
+ attMap -= attMap.min()
15
+ attMap /= attMap.max()
16
+ cmap = plt.get_cmap("jet")
17
+ attMapV = cmap(attMap)
18
+ attMapV = np.delete(attMapV, 3, 2)
19
+ if overlap:
20
+ attMap = (
21
+ 1 * (1 - attMap**0.7).reshape(attMap.shape + (1,)) * img
22
+ + (attMap**0.7).reshape(attMap.shape + (1,)) * attMapV
23
+ )
24
+ return attMap
unimernet/common/logger.py ADDED
@@ -0,0 +1,195 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ import datetime
9
+ import logging
10
+ import time
11
+ from collections import defaultdict, deque
12
+
13
+ import torch
14
+ import torch.distributed as dist
15
+
16
+ from unimernet.common import dist_utils
17
+
18
+
19
+ class SmoothedValue(object):
20
+ """Track a series of values and provide access to smoothed values over a
21
+ window or the global series average.
22
+ """
23
+
24
+ def __init__(self, window_size=20, fmt=None):
25
+ if fmt is None:
26
+ fmt = "{median:.4f} ({global_avg:.4f})"
27
+ self.deque = deque(maxlen=window_size)
28
+ self.total = 0.0
29
+ self.count = 0
30
+ self.fmt = fmt
31
+
32
+ def update(self, value, n=1):
33
+ self.deque.append(value)
34
+ self.count += n
35
+ self.total += value * n
36
+
37
+ def synchronize_between_processes(self):
38
+ """
39
+ Warning: does not synchronize the deque!
40
+ """
41
+ if not dist_utils.is_dist_avail_and_initialized():
42
+ return
43
+ t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
44
+ dist.barrier()
45
+ dist.all_reduce(t)
46
+ t = t.tolist()
47
+ self.count = int(t[0])
48
+ self.total = t[1]
49
+
50
+ @property
51
+ def median(self):
52
+ d = torch.tensor(list(self.deque))
53
+ return d.median().item()
54
+
55
+ @property
56
+ def avg(self):
57
+ d = torch.tensor(list(self.deque), dtype=torch.float32)
58
+ return d.mean().item()
59
+
60
+ @property
61
+ def global_avg(self):
62
+ return self.total / self.count
63
+
64
+ @property
65
+ def max(self):
66
+ return max(self.deque)
67
+
68
+ @property
69
+ def value(self):
70
+ return self.deque[-1]
71
+
72
+ def __str__(self):
73
+ return self.fmt.format(
74
+ median=self.median,
75
+ avg=self.avg,
76
+ global_avg=self.global_avg,
77
+ max=self.max,
78
+ value=self.value,
79
+ )
80
+
81
+
82
+ class MetricLogger(object):
83
+ def __init__(self, delimiter="\t"):
84
+ self.meters = defaultdict(SmoothedValue)
85
+ self.delimiter = delimiter
86
+
87
+ def update(self, **kwargs):
88
+ for k, v in kwargs.items():
89
+ if isinstance(v, torch.Tensor):
90
+ v = v.item()
91
+ assert isinstance(v, (float, int))
92
+ self.meters[k].update(v)
93
+
94
+ def __getattr__(self, attr):
95
+ if attr in self.meters:
96
+ return self.meters[attr]
97
+ if attr in self.__dict__:
98
+ return self.__dict__[attr]
99
+ raise AttributeError(
100
+ "'{}' object has no attribute '{}'".format(type(self).__name__, attr)
101
+ )
102
+
103
+ def __str__(self):
104
+ loss_str = []
105
+ for name, meter in self.meters.items():
106
+ loss_str.append("{}: {}".format(name, str(meter)))
107
+ return self.delimiter.join(loss_str)
108
+
109
+ def global_avg(self):
110
+ loss_str = []
111
+ for name, meter in self.meters.items():
112
+ loss_str.append("{}: {:.4f}".format(name, meter.global_avg))
113
+ return self.delimiter.join(loss_str)
114
+
115
+ def synchronize_between_processes(self):
116
+ for meter in self.meters.values():
117
+ meter.synchronize_between_processes()
118
+
119
+ def add_meter(self, name, meter):
120
+ self.meters[name] = meter
121
+
122
+ def log_every(self, iterable, print_freq, header=None):
123
+ i = 0
124
+ if not header:
125
+ header = ""
126
+ start_time = time.time()
127
+ end = time.time()
128
+ iter_time = SmoothedValue(fmt="{avg:.4f}")
129
+ data_time = SmoothedValue(fmt="{avg:.4f}")
130
+ space_fmt = ":" + str(len(str(len(iterable)))) + "d"
131
+ log_msg = [
132
+ header,
133
+ "[{0" + space_fmt + "}/{1}]",
134
+ "eta: {eta}",
135
+ "{meters}",
136
+ "time: {time}",
137
+ "data: {data}",
138
+ ]
139
+ if torch.cuda.is_available():
140
+ log_msg.append("max mem: {memory:.0f}")
141
+ log_msg = self.delimiter.join(log_msg)
142
+ MB = 1024.0 * 1024.0
143
+ for obj in iterable:
144
+ data_time.update(time.time() - end)
145
+ yield obj
146
+ iter_time.update(time.time() - end)
147
+ if i % print_freq == 0 or i == len(iterable) - 1:
148
+ eta_seconds = iter_time.global_avg * (len(iterable) - i)
149
+ eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
150
+ if torch.cuda.is_available():
151
+ print(
152
+ log_msg.format(
153
+ i,
154
+ len(iterable),
155
+ eta=eta_string,
156
+ meters=str(self),
157
+ time=str(iter_time),
158
+ data=str(data_time),
159
+ memory=torch.cuda.max_memory_allocated() / MB,
160
+ )
161
+ )
162
+ else:
163
+ print(
164
+ log_msg.format(
165
+ i,
166
+ len(iterable),
167
+ eta=eta_string,
168
+ meters=str(self),
169
+ time=str(iter_time),
170
+ data=str(data_time),
171
+ )
172
+ )
173
+ i += 1
174
+ end = time.time()
175
+ total_time = time.time() - start_time
176
+ total_time_str = str(datetime.timedelta(seconds=int(total_time)))
177
+ print(
178
+ "{} Total time: {} ({:.4f} s / it)".format(
179
+ header, total_time_str, total_time / len(iterable)
180
+ )
181
+ )
182
+
183
+
184
+ class AttrDict(dict):
185
+ def __init__(self, *args, **kwargs):
186
+ super(AttrDict, self).__init__(*args, **kwargs)
187
+ self.__dict__ = self
188
+
189
+
190
+ def setup_logger():
191
+ logging.basicConfig(
192
+ level=logging.INFO if dist_utils.is_main_process() else logging.WARN,
193
+ format="%(asctime)s [%(levelname)s] %(message)s",
194
+ handlers=[logging.StreamHandler()],
195
+ )
unimernet/common/optims.py ADDED
@@ -0,0 +1,120 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ import math
9
+
10
+ from unimernet.common.registry import registry
11
+
12
+
13
+ @registry.register_lr_scheduler("linear_warmup_step_lr")
14
+ class LinearWarmupStepLRScheduler:
15
+ def __init__(
16
+ self,
17
+ optimizer,
18
+ max_epoch,
19
+ min_lr,
20
+ init_lr,
21
+ decay_rate=1,
22
+ warmup_start_lr=-1,
23
+ warmup_steps=0,
24
+ **kwargs
25
+ ):
26
+ self.optimizer = optimizer
27
+
28
+ self.max_epoch = max_epoch
29
+ self.min_lr = min_lr
30
+
31
+ self.decay_rate = decay_rate
32
+
33
+ self.init_lr = init_lr
34
+ self.warmup_steps = warmup_steps
35
+ self.warmup_start_lr = warmup_start_lr if warmup_start_lr >= 0 else init_lr
36
+
37
+ def step(self, cur_epoch, cur_step):
38
+ if cur_epoch == 0:
39
+ warmup_lr_schedule(
40
+ step=cur_step,
41
+ optimizer=self.optimizer,
42
+ max_step=self.warmup_steps,
43
+ init_lr=self.warmup_start_lr,
44
+ max_lr=self.init_lr,
45
+ )
46
+ else:
47
+ step_lr_schedule(
48
+ epoch=cur_epoch,
49
+ optimizer=self.optimizer,
50
+ init_lr=self.init_lr,
51
+ min_lr=self.min_lr,
52
+ decay_rate=self.decay_rate,
53
+ )
54
+
55
+
56
+ @registry.register_lr_scheduler("linear_warmup_cosine_lr")
57
+ class LinearWarmupCosineLRScheduler:
58
+ def __init__(
59
+ self,
60
+ optimizer,
61
+ max_epoch,
62
+ min_lr,
63
+ init_lr,
64
+ iters_per_epoch,
65
+ warmup_steps=0,
66
+ warmup_start_lr=-1,
67
+ **kwargs
68
+ ):
69
+ self.optimizer = optimizer
70
+
71
+ self.max_epoch = max_epoch
72
+ self.min_lr = min_lr
73
+
74
+ self.init_lr = init_lr
75
+ self.warmup_steps = warmup_steps
76
+ self.iters_per_epoch = iters_per_epoch
77
+ self.warmup_start_lr = warmup_start_lr if warmup_start_lr >= 0 else init_lr
78
+
79
+ def step(self, cur_epoch, cur_step):
80
+ # assuming the warmup iters less than one epoch
81
+ total_steps = cur_epoch * self.iters_per_epoch + cur_step
82
+ if total_steps < self.warmup_steps:
83
+ warmup_lr_schedule(
84
+ step=cur_step,
85
+ optimizer=self.optimizer,
86
+ max_step=self.warmup_steps,
87
+ init_lr=self.warmup_start_lr,
88
+ max_lr=self.init_lr,
89
+ )
90
+ else:
91
+ cosine_lr_schedule(
92
+ epoch=total_steps,
93
+ optimizer=self.optimizer,
94
+ max_epoch=self.max_epoch * self.iters_per_epoch,
95
+ init_lr=self.init_lr,
96
+ min_lr=self.min_lr,
97
+ )
98
+
99
+
100
+ def cosine_lr_schedule(optimizer, epoch, max_epoch, init_lr, min_lr):
101
+ """Decay the learning rate"""
102
+ lr = (init_lr - min_lr) * 0.5 * (
103
+ 1.0 + math.cos(math.pi * epoch / max_epoch)
104
+ ) + min_lr
105
+ for param_group in optimizer.param_groups:
106
+ param_group["lr"] = lr
107
+
108
+
109
+ def warmup_lr_schedule(optimizer, step, max_step, init_lr, max_lr):
110
+ """Warmup the learning rate"""
111
+ lr = min(max_lr, init_lr + (max_lr - init_lr) * step / max(max_step, 1))
112
+ for param_group in optimizer.param_groups:
113
+ param_group["lr"] = lr
114
+
115
+
116
+ def step_lr_schedule(optimizer, epoch, init_lr, min_lr, decay_rate):
117
+ """Decay the learning rate"""
118
+ lr = max(min_lr, init_lr * (decay_rate**epoch))
119
+ for param_group in optimizer.param_groups:
120
+ param_group["lr"] = lr
unimernet/common/registry.py ADDED
@@ -0,0 +1,329 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+
9
+ class Registry:
10
+ mapping = {
11
+ "builder_name_mapping": {},
12
+ "task_name_mapping": {},
13
+ "processor_name_mapping": {},
14
+ "model_name_mapping": {},
15
+ "lr_scheduler_name_mapping": {},
16
+ "runner_name_mapping": {},
17
+ "state": {},
18
+ "paths": {},
19
+ }
20
+
21
+ @classmethod
22
+ def register_builder(cls, name):
23
+ r"""Register a dataset builder to registry with key 'name'
24
+
25
+ Args:
26
+ name: Key with which the builder will be registered.
27
+
28
+ Usage:
29
+
30
+ from unimernet.common.registry import registry
31
+ from unimernet.datasets.base_dataset_builder import BaseDatasetBuilder
32
+ """
33
+
34
+ def wrap(builder_cls):
35
+ from unimernet.datasets.builders.base_dataset_builder import BaseDatasetBuilder
36
+
37
+ assert issubclass(
38
+ builder_cls, BaseDatasetBuilder
39
+ ), "All builders must inherit BaseDatasetBuilder class, found {}".format(
40
+ builder_cls
41
+ )
42
+ if name in cls.mapping["builder_name_mapping"]:
43
+ raise KeyError(
44
+ "Name '{}' already registered for {}.".format(
45
+ name, cls.mapping["builder_name_mapping"][name]
46
+ )
47
+ )
48
+ cls.mapping["builder_name_mapping"][name] = builder_cls
49
+ return builder_cls
50
+
51
+ return wrap
52
+
53
+ @classmethod
54
+ def register_task(cls, name):
55
+ r"""Register a task to registry with key 'name'
56
+
57
+ Args:
58
+ name: Key with which the task will be registered.
59
+
60
+ Usage:
61
+
62
+ from unimernet.common.registry import registry
63
+ """
64
+
65
+ def wrap(task_cls):
66
+ from unimernet.tasks.base_task import BaseTask
67
+
68
+ assert issubclass(
69
+ task_cls, BaseTask
70
+ ), "All tasks must inherit BaseTask class"
71
+ if name in cls.mapping["task_name_mapping"]:
72
+ raise KeyError(
73
+ "Name '{}' already registered for {}.".format(
74
+ name, cls.mapping["task_name_mapping"][name]
75
+ )
76
+ )
77
+ cls.mapping["task_name_mapping"][name] = task_cls
78
+ return task_cls
79
+
80
+ return wrap
81
+
82
+ @classmethod
83
+ def register_model(cls, name):
84
+ r"""Register a task to registry with key 'name'
85
+
86
+ Args:
87
+ name: Key with which the task will be registered.
88
+
89
+ Usage:
90
+
91
+ from unimernet.common.registry import registry
92
+ """
93
+
94
+ def wrap(model_cls):
95
+ from unimernet.models import BaseModel
96
+
97
+ assert issubclass(
98
+ model_cls, BaseModel
99
+ ), "All models must inherit BaseModel class"
100
+ if name in cls.mapping["model_name_mapping"]:
101
+ raise KeyError(
102
+ "Name '{}' already registered for {}.".format(
103
+ name, cls.mapping["model_name_mapping"][name]
104
+ )
105
+ )
106
+ cls.mapping["model_name_mapping"][name] = model_cls
107
+ return model_cls
108
+
109
+ return wrap
110
+
111
+ @classmethod
112
+ def register_processor(cls, name):
113
+ r"""Register a processor to registry with key 'name'
114
+
115
+ Args:
116
+ name: Key with which the task will be registered.
117
+
118
+ Usage:
119
+
120
+ from unimernet.common.registry import registry
121
+ """
122
+
123
+ def wrap(processor_cls):
124
+ from unimernet.processors import BaseProcessor
125
+
126
+ assert issubclass(
127
+ processor_cls, BaseProcessor
128
+ ), "All processors must inherit BaseProcessor class"
129
+ if name in cls.mapping["processor_name_mapping"]:
130
+ raise KeyError(
131
+ "Name '{}' already registered for {}.".format(
132
+ name, cls.mapping["processor_name_mapping"][name]
133
+ )
134
+ )
135
+ cls.mapping["processor_name_mapping"][name] = processor_cls
136
+ return processor_cls
137
+
138
+ return wrap
139
+
140
+ @classmethod
141
+ def register_lr_scheduler(cls, name):
142
+ r"""Register a model to registry with key 'name'
143
+
144
+ Args:
145
+ name: Key with which the task will be registered.
146
+
147
+ Usage:
148
+
149
+ from unimernet.common.registry import registry
150
+ """
151
+
152
+ def wrap(lr_sched_cls):
153
+ if name in cls.mapping["lr_scheduler_name_mapping"]:
154
+ raise KeyError(
155
+ "Name '{}' already registered for {}.".format(
156
+ name, cls.mapping["lr_scheduler_name_mapping"][name]
157
+ )
158
+ )
159
+ cls.mapping["lr_scheduler_name_mapping"][name] = lr_sched_cls
160
+ return lr_sched_cls
161
+
162
+ return wrap
163
+
164
+ @classmethod
165
+ def register_runner(cls, name):
166
+ r"""Register a model to registry with key 'name'
167
+
168
+ Args:
169
+ name: Key with which the task will be registered.
170
+
171
+ Usage:
172
+
173
+ from unimernet.common.registry import registry
174
+ """
175
+
176
+ def wrap(runner_cls):
177
+ if name in cls.mapping["runner_name_mapping"]:
178
+ raise KeyError(
179
+ "Name '{}' already registered for {}.".format(
180
+ name, cls.mapping["runner_name_mapping"][name]
181
+ )
182
+ )
183
+ cls.mapping["runner_name_mapping"][name] = runner_cls
184
+ return runner_cls
185
+
186
+ return wrap
187
+
188
+ @classmethod
189
+ def register_path(cls, name, path):
190
+ r"""Register a path to registry with key 'name'
191
+
192
+ Args:
193
+ name: Key with which the path will be registered.
194
+
195
+ Usage:
196
+
197
+ from unimernet.common.registry import registry
198
+ """
199
+ assert isinstance(path, str), "All path must be str."
200
+ if name in cls.mapping["paths"]:
201
+ raise KeyError("Name '{}' already registered.".format(name))
202
+ cls.mapping["paths"][name] = path
203
+
204
+ @classmethod
205
+ def register(cls, name, obj):
206
+ r"""Register an item to registry with key 'name'
207
+
208
+ Args:
209
+ name: Key with which the item will be registered.
210
+
211
+ Usage::
212
+
213
+ from unimernet.common.registry import registry
214
+
215
+ registry.register("config", {})
216
+ """
217
+ path = name.split(".")
218
+ current = cls.mapping["state"]
219
+
220
+ for part in path[:-1]:
221
+ if part not in current:
222
+ current[part] = {}
223
+ current = current[part]
224
+
225
+ current[path[-1]] = obj
226
+
227
+ # @classmethod
228
+ # def get_trainer_class(cls, name):
229
+ # return cls.mapping["trainer_name_mapping"].get(name, None)
230
+
231
+ @classmethod
232
+ def get_builder_class(cls, name):
233
+ return cls.mapping["builder_name_mapping"].get(name, None)
234
+
235
+ @classmethod
236
+ def get_model_class(cls, name):
237
+ return cls.mapping["model_name_mapping"].get(name, None)
238
+
239
+ @classmethod
240
+ def get_task_class(cls, name):
241
+ return cls.mapping["task_name_mapping"].get(name, None)
242
+
243
+ @classmethod
244
+ def get_processor_class(cls, name):
245
+ return cls.mapping["processor_name_mapping"].get(name, None)
246
+
247
+ @classmethod
248
+ def get_lr_scheduler_class(cls, name):
249
+ return cls.mapping["lr_scheduler_name_mapping"].get(name, None)
250
+
251
+ @classmethod
252
+ def get_runner_class(cls, name):
253
+ return cls.mapping["runner_name_mapping"].get(name, None)
254
+
255
+ @classmethod
256
+ def list_runners(cls):
257
+ return sorted(cls.mapping["runner_name_mapping"].keys())
258
+
259
+ @classmethod
260
+ def list_models(cls):
261
+ return sorted(cls.mapping["model_name_mapping"].keys())
262
+
263
+ @classmethod
264
+ def list_tasks(cls):
265
+ return sorted(cls.mapping["task_name_mapping"].keys())
266
+
267
+ @classmethod
268
+ def list_processors(cls):
269
+ return sorted(cls.mapping["processor_name_mapping"].keys())
270
+
271
+ @classmethod
272
+ def list_lr_schedulers(cls):
273
+ return sorted(cls.mapping["lr_scheduler_name_mapping"].keys())
274
+
275
+ @classmethod
276
+ def list_datasets(cls):
277
+ return sorted(cls.mapping["builder_name_mapping"].keys())
278
+
279
+ @classmethod
280
+ def get_path(cls, name):
281
+ return cls.mapping["paths"].get(name, None)
282
+
283
+ @classmethod
284
+ def get(cls, name, default=None, no_warning=False):
285
+ r"""Get an item from registry with key 'name'
286
+
287
+ Args:
288
+ name (string): Key whose value needs to be retrieved.
289
+ default: If passed and key is not in registry, default value will
290
+ be returned with a warning. Default: None
291
+ no_warning (bool): If passed as True, warning when key doesn't exist
292
+ will not be generated. Useful for MMF's
293
+ internal operations. Default: False
294
+ """
295
+ original_name = name
296
+ name = name.split(".")
297
+ value = cls.mapping["state"]
298
+ for subname in name:
299
+ value = value.get(subname, default)
300
+ if value is default:
301
+ break
302
+
303
+ if (
304
+ "writer" in cls.mapping["state"]
305
+ and value == default
306
+ and no_warning is False
307
+ ):
308
+ cls.mapping["state"]["writer"].warning(
309
+ "Key {} is not present in registry, returning default value "
310
+ "of {}".format(original_name, default)
311
+ )
312
+ return value
313
+
314
+ @classmethod
315
+ def unregister(cls, name):
316
+ r"""Remove an item from registry with key 'name'
317
+
318
+ Args:
319
+ name: Key which needs to be removed.
320
+ Usage::
321
+
322
+ from mmf.common.registry import registry
323
+
324
+ config = registry.unregister("config")
325
+ """
326
+ return cls.mapping["state"].pop(name, None)
327
+
328
+
329
+ registry = Registry()
unimernet/common/utils.py ADDED
@@ -0,0 +1,424 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ import io
9
+ import json
10
+ import logging
11
+ import os
12
+ import pickle
13
+ import re
14
+ import shutil
15
+ import urllib
16
+ import urllib.error
17
+ import urllib.request
18
+ from typing import Optional
19
+ from urllib.parse import urlparse
20
+
21
+ import numpy as np
22
+ import pandas as pd
23
+ import yaml
24
+ from iopath.common.download import download
25
+ from iopath.common.file_io import file_lock, g_pathmgr
26
+ from unimernet.common.registry import registry
27
+ from torch.utils.model_zoo import tqdm
28
+ from torchvision.datasets.utils import (
29
+ check_integrity,
30
+ download_file_from_google_drive,
31
+ extract_archive,
32
+ )
33
+
34
+
35
+ def now():
36
+ from datetime import datetime
37
+
38
+ return datetime.now().strftime("%Y%m%d%H%M")[:-1]
39
+
40
+
41
+ def is_url(url_or_filename):
42
+ parsed = urlparse(url_or_filename)
43
+ return parsed.scheme in ("http", "https")
44
+
45
+
46
+ def get_cache_path(rel_path):
47
+ return os.path.expanduser(os.path.join(registry.get_path("cache_root"), rel_path))
48
+
49
+
50
+ def get_abs_path(rel_path):
51
+ return os.path.join(registry.get_path("library_root"), rel_path)
52
+
53
+
54
+ def load_json(filename):
55
+ with open(filename, "r") as f:
56
+ return json.load(f)
57
+
58
+
59
+ # The following are adapted from torchvision and vissl
60
+ # torchvision: https://github.com/pytorch/vision
61
+ # vissl: https://github.com/facebookresearch/vissl/blob/main/vissl/utils/download.py
62
+
63
+
64
+ def makedir(dir_path):
65
+ """
66
+ Create the directory if it does not exist.
67
+ """
68
+ is_success = False
69
+ try:
70
+ if not g_pathmgr.exists(dir_path):
71
+ g_pathmgr.mkdirs(dir_path)
72
+ is_success = True
73
+ except BaseException:
74
+ print(f"Error creating directory: {dir_path}")
75
+ return is_success
76
+
77
+
78
+ def get_redirected_url(url: str):
79
+ """
80
+ Given a URL, returns the URL it redirects to or the
81
+ original URL in case of no indirection
82
+ """
83
+ import requests
84
+
85
+ with requests.Session() as session:
86
+ with session.get(url, stream=True, allow_redirects=True) as response:
87
+ if response.history:
88
+ return response.url
89
+ else:
90
+ return url
91
+
92
+
93
+ def to_google_drive_download_url(view_url: str) -> str:
94
+ """
95
+ Utility function to transform a view URL of google drive
96
+ to a download URL for google drive
97
+ Example input:
98
+ https://drive.google.com/file/d/137RyRjvTBkBiIfeYBNZBtViDHQ6_Ewsp/view
99
+ Example output:
100
+ https://drive.google.com/uc?export=download&id=137RyRjvTBkBiIfeYBNZBtViDHQ6_Ewsp
101
+ """
102
+ splits = view_url.split("/")
103
+ assert splits[-1] == "view"
104
+ file_id = splits[-2]
105
+ return f"https://drive.google.com/uc?export=download&id={file_id}"
106
+
107
+
108
+ def download_google_drive_url(url: str, output_path: str, output_file_name: str):
109
+ """
110
+ Download a file from google drive
111
+ Downloading an URL from google drive requires confirmation when
112
+ the file of the size is too big (google drive notifies that
113
+ anti-viral checks cannot be performed on such files)
114
+ """
115
+ import requests
116
+
117
+ with requests.Session() as session:
118
+
119
+ # First get the confirmation token and append it to the URL
120
+ with session.get(url, stream=True, allow_redirects=True) as response:
121
+ for k, v in response.cookies.items():
122
+ if k.startswith("download_warning"):
123
+ url = url + "&confirm=" + v
124
+
125
+ # Then download the content of the file
126
+ with session.get(url, stream=True, verify=True) as response:
127
+ makedir(output_path)
128
+ path = os.path.join(output_path, output_file_name)
129
+ total_size = int(response.headers.get("Content-length", 0))
130
+ with open(path, "wb") as file:
131
+ from tqdm import tqdm
132
+
133
+ with tqdm(total=total_size) as progress_bar:
134
+ for block in response.iter_content(
135
+ chunk_size=io.DEFAULT_BUFFER_SIZE
136
+ ):
137
+ file.write(block)
138
+ progress_bar.update(len(block))
139
+
140
+
141
+ def _get_google_drive_file_id(url: str) -> Optional[str]:
142
+ parts = urlparse(url)
143
+
144
+ if re.match(r"(drive|docs)[.]google[.]com", parts.netloc) is None:
145
+ return None
146
+
147
+ match = re.match(r"/file/d/(?P<id>[^/]*)", parts.path)
148
+ if match is None:
149
+ return None
150
+
151
+ return match.group("id")
152
+
153
+
154
+ def _urlretrieve(url: str, filename: str, chunk_size: int = 1024) -> None:
155
+ with open(filename, "wb") as fh:
156
+ with urllib.request.urlopen(
157
+ urllib.request.Request(url, headers={"User-Agent": "vissl"})
158
+ ) as response:
159
+ with tqdm(total=response.length) as pbar:
160
+ for chunk in iter(lambda: response.read(chunk_size), ""):
161
+ if not chunk:
162
+ break
163
+ pbar.update(chunk_size)
164
+ fh.write(chunk)
165
+
166
+
167
+ def download_url(
168
+ url: str,
169
+ root: str,
170
+ filename: Optional[str] = None,
171
+ md5: Optional[str] = None,
172
+ ) -> None:
173
+ """Download a file from a url and place it in root.
174
+ Args:
175
+ url (str): URL to download file from
176
+ root (str): Directory to place downloaded file in
177
+ filename (str, optional): Name to save the file under.
178
+ If None, use the basename of the URL.
179
+ md5 (str, optional): MD5 checksum of the download. If None, do not check
180
+ """
181
+ root = os.path.expanduser(root)
182
+ if not filename:
183
+ filename = os.path.basename(url)
184
+ fpath = os.path.join(root, filename)
185
+
186
+ makedir(root)
187
+
188
+ # check if file is already present locally
189
+ if check_integrity(fpath, md5):
190
+ print("Using downloaded and verified file: " + fpath)
191
+ return
192
+
193
+ # expand redirect chain if needed
194
+ url = get_redirected_url(url)
195
+
196
+ # check if file is located on Google Drive
197
+ file_id = _get_google_drive_file_id(url)
198
+ if file_id is not None:
199
+ return download_file_from_google_drive(file_id, root, filename, md5)
200
+
201
+ # download the file
202
+ try:
203
+ print("Downloading " + url + " to " + fpath)
204
+ _urlretrieve(url, fpath)
205
+ except (urllib.error.URLError, IOError) as e: # type: ignore[attr-defined]
206
+ if url[:5] == "https":
207
+ url = url.replace("https:", "http:")
208
+ print(
209
+ "Failed download. Trying https -> http instead."
210
+ " Downloading " + url + " to " + fpath
211
+ )
212
+ _urlretrieve(url, fpath)
213
+ else:
214
+ raise e
215
+
216
+ # check integrity of downloaded file
217
+ if not check_integrity(fpath, md5):
218
+ raise RuntimeError("File not found or corrupted.")
219
+
220
+
221
+ def download_and_extract_archive(
222
+ url: str,
223
+ download_root: str,
224
+ extract_root: Optional[str] = None,
225
+ filename: Optional[str] = None,
226
+ md5: Optional[str] = None,
227
+ remove_finished: bool = False,
228
+ ) -> None:
229
+ download_root = os.path.expanduser(download_root)
230
+ if extract_root is None:
231
+ extract_root = download_root
232
+ if not filename:
233
+ filename = os.path.basename(url)
234
+
235
+ download_url(url, download_root, filename, md5)
236
+
237
+ archive = os.path.join(download_root, filename)
238
+ print("Extracting {} to {}".format(archive, extract_root))
239
+ extract_archive(archive, extract_root, remove_finished)
240
+
241
+
242
+ def cache_url(url: str, cache_dir: str) -> str:
243
+ """
244
+ This implementation downloads the remote resource and caches it locally.
245
+ The resource will only be downloaded if not previously requested.
246
+ """
247
+ parsed_url = urlparse(url)
248
+ dirname = os.path.join(cache_dir, os.path.dirname(parsed_url.path.lstrip("/")))
249
+ makedir(dirname)
250
+ filename = url.split("/")[-1]
251
+ cached = os.path.join(dirname, filename)
252
+ with file_lock(cached):
253
+ if not os.path.isfile(cached):
254
+ logging.info(f"Downloading {url} to {cached} ...")
255
+ cached = download(url, dirname, filename=filename)
256
+ logging.info(f"URL {url} cached in {cached}")
257
+ return cached
258
+
259
+
260
+ # TODO (prigoyal): convert this into RAII-style API
261
+ def create_file_symlink(file1, file2):
262
+ """
263
+ Simply create the symlinks for a given file1 to file2.
264
+ Useful during model checkpointing to symlinks to the
265
+ latest successful checkpoint.
266
+ """
267
+ try:
268
+ if g_pathmgr.exists(file2):
269
+ g_pathmgr.rm(file2)
270
+ g_pathmgr.symlink(file1, file2)
271
+ except Exception as e:
272
+ logging.info(f"Could NOT create symlink. Error: {e}")
273
+
274
+
275
+ def save_file(data, filename, append_to_json=True, verbose=True):
276
+ """
277
+ Common i/o utility to handle saving data to various file formats.
278
+ Supported:
279
+ .pkl, .pickle, .npy, .json
280
+ Specifically for .json, users have the option to either append (default)
281
+ or rewrite by passing in Boolean value to append_to_json.
282
+ """
283
+ if verbose:
284
+ logging.info(f"Saving data to file: {filename}")
285
+ file_ext = os.path.splitext(filename)[1]
286
+ if file_ext in [".pkl", ".pickle"]:
287
+ with g_pathmgr.open(filename, "wb") as fopen:
288
+ pickle.dump(data, fopen, pickle.HIGHEST_PROTOCOL)
289
+ elif file_ext == ".npy":
290
+ with g_pathmgr.open(filename, "wb") as fopen:
291
+ np.save(fopen, data)
292
+ elif file_ext == ".json":
293
+ if append_to_json:
294
+ with g_pathmgr.open(filename, "a") as fopen:
295
+ fopen.write(json.dumps(data, sort_keys=True) + "\n")
296
+ fopen.flush()
297
+ else:
298
+ with g_pathmgr.open(filename, "w") as fopen:
299
+ fopen.write(json.dumps(data, sort_keys=True) + "\n")
300
+ fopen.flush()
301
+ elif file_ext == ".yaml":
302
+ with g_pathmgr.open(filename, "w") as fopen:
303
+ dump = yaml.dump(data)
304
+ fopen.write(dump)
305
+ fopen.flush()
306
+ else:
307
+ raise Exception(f"Saving {file_ext} is not supported yet")
308
+
309
+ if verbose:
310
+ logging.info(f"Saved data to file: {filename}")
311
+
312
+
313
+ def load_file(filename, mmap_mode=None, verbose=True, allow_pickle=False):
314
+ """
315
+ Common i/o utility to handle loading data from various file formats.
316
+ Supported:
317
+ .pkl, .pickle, .npy, .json
318
+ For the npy files, we support reading the files in mmap_mode.
319
+ If the mmap_mode of reading is not successful, we load data without the
320
+ mmap_mode.
321
+ """
322
+ if verbose:
323
+ logging.info(f"Loading data from file: {filename}")
324
+
325
+ file_ext = os.path.splitext(filename)[1]
326
+ if file_ext == ".txt":
327
+ with g_pathmgr.open(filename, "r") as fopen:
328
+ data = fopen.readlines()
329
+ elif file_ext in [".pkl", ".pickle"]:
330
+ with g_pathmgr.open(filename, "rb") as fopen:
331
+ data = pickle.load(fopen, encoding="latin1")
332
+ elif file_ext == ".npy":
333
+ if mmap_mode:
334
+ try:
335
+ with g_pathmgr.open(filename, "rb") as fopen:
336
+ data = np.load(
337
+ fopen,
338
+ allow_pickle=allow_pickle,
339
+ encoding="latin1",
340
+ mmap_mode=mmap_mode,
341
+ )
342
+ except ValueError as e:
343
+ logging.info(
344
+ f"Could not mmap {filename}: {e}. Trying without g_pathmgr"
345
+ )
346
+ data = np.load(
347
+ filename,
348
+ allow_pickle=allow_pickle,
349
+ encoding="latin1",
350
+ mmap_mode=mmap_mode,
351
+ )
352
+ logging.info("Successfully loaded without g_pathmgr")
353
+ except Exception:
354
+ logging.info("Could not mmap without g_pathmgr. Trying without mmap")
355
+ with g_pathmgr.open(filename, "rb") as fopen:
356
+ data = np.load(fopen, allow_pickle=allow_pickle, encoding="latin1")
357
+ else:
358
+ with g_pathmgr.open(filename, "rb") as fopen:
359
+ data = np.load(fopen, allow_pickle=allow_pickle, encoding="latin1")
360
+ elif file_ext == ".json":
361
+ with g_pathmgr.open(filename, "r") as fopen:
362
+ data = json.load(fopen)
363
+ elif file_ext == ".yaml":
364
+ with g_pathmgr.open(filename, "r") as fopen:
365
+ data = yaml.load(fopen, Loader=yaml.FullLoader)
366
+ elif file_ext == ".csv":
367
+ with g_pathmgr.open(filename, "r") as fopen:
368
+ data = pd.read_csv(fopen)
369
+ else:
370
+ raise Exception(f"Reading from {file_ext} is not supported yet")
371
+ return data
372
+
373
+
374
+ def abspath(resource_path: str):
375
+ """
376
+ Make a path absolute, but take into account prefixes like
377
+ "http://" or "manifold://"
378
+ """
379
+ regex = re.compile(r"^\w+://")
380
+ if regex.match(resource_path) is None:
381
+ return os.path.abspath(resource_path)
382
+ else:
383
+ return resource_path
384
+
385
+
386
+ def makedir(dir_path):
387
+ """
388
+ Create the directory if it does not exist.
389
+ """
390
+ is_success = False
391
+ try:
392
+ if not g_pathmgr.exists(dir_path):
393
+ g_pathmgr.mkdirs(dir_path)
394
+ is_success = True
395
+ except BaseException:
396
+ logging.info(f"Error creating directory: {dir_path}")
397
+ return is_success
398
+
399
+
400
+ def is_url(input_url):
401
+ """
402
+ Check if an input string is a url. look for http(s):// and ignoring the case
403
+ """
404
+ is_url = re.match(r"^(?:http)s?://", input_url, re.IGNORECASE) is not None
405
+ return is_url
406
+
407
+
408
+ def cleanup_dir(dir):
409
+ """
410
+ Utility for deleting a directory. Useful for cleaning the storage space
411
+ that contains various training artifacts like checkpoints, data etc.
412
+ """
413
+ if os.path.exists(dir):
414
+ logging.info(f"Deleting directory: {dir}")
415
+ shutil.rmtree(dir)
416
+ logging.info(f"Deleted contents of directory: {dir}")
417
+
418
+
419
+ def get_file_size(filename):
420
+ """
421
+ Given a file, get the size of file in MB
422
+ """
423
+ size_in_mb = os.path.getsize(filename) / float(1024**2)
424
+ return size_in_mb
unimernet/configs/datasets/formula/formula_eval.yaml ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ datasets:
2
+ formula_rec_eval:
3
+ data_type: images
4
+ build_info:
5
+ images: /mnt/petrelfs/share_data/hanxiao/latex-ocr/pdf/val
6
+ annotation: /mnt/petrelfs/share_data/hanxiao/latex-ocr/pdf/pdfmath.txt
unimernet/configs/datasets/formula/formula_train.yaml ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ datasets:
2
+ formula_rec_train:
3
+ data_type: images
4
+ build_info:
5
+ images: /mnt/petrelfs/share_data/hanxiao/latex-ocr/pdf/train
6
+ annotation: /mnt/petrelfs/share_data/hanxiao/latex-ocr/pdf/pdfmath.txt
unimernet/configs/datasets/formula/multi_scale_formula_train.yaml ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ datasets:
2
+ multi_scale_formula_rec_train:
3
+ data_type: images
4
+ build_info:
5
+ images: /mnt/petrelfs/share_data/hanxiao/latex-ocr/pdf/train
6
+ annotation: /mnt/petrelfs/share_data/hanxiao/latex-ocr/pdf/pdfmath.txt
7
+
8
+ vis_processor:
9
+ train:
10
+ name: "formula_image_multi_scale_train"
11
+ all_scales:
12
+ - [ 96, 336 ]
13
+ - [ 128, 448 ]
14
+ - [ 192, 672 ]
15
+ - [ 288, 1008 ]
16
+ - [ 384, 1344 ]
17
+
18
+ text_processor:
19
+ train:
20
+ name: "blip_caption"
21
+ max_words: 256
unimernet/configs/default.yaml ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2022, salesforce.com, inc.
2
+ # All rights reserved.
3
+ # SPDX-License-Identifier: BSD-3-Clause
4
+ # For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
5
+
6
+ env:
7
+ # For default users
8
+ # cache_root: "cache"
9
+ # For internal use with persistent storage
10
+ cache_root: "/export/home/.cache/vigc"
unimernet/configs/models/unimernet_base.yaml ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ model:
2
+ arch: unimernet
3
+ load_finetuned: False
4
+ load_pretrained: False
5
+ pretrained: "path/to/pretrained/weight"
6
+ finetuned: ""
7
+ tokenizer_name: nougat
8
+ tokenizer_config:
9
+ path: ./models/unimernet
10
+ model_name: unimernet
11
+ model_config:
12
+ max_seq_len: 384
13
+
14
+
15
+ preprocess:
16
+ vis_processor:
17
+ train:
18
+ name: "formula_image_train"
19
+ image_size:
20
+ - 192
21
+ - 672
22
+ eval:
23
+ name: "formula_image_eval"
24
+ image_size:
25
+ - 192
26
+ - 672
27
+ text_processor:
28
+ train:
29
+ name: "blip_caption"
30
+ eval:
31
+ name: "blip_caption"
unimernet/datasets/__init__.py ADDED
File without changes
unimernet/datasets/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (162 Bytes). View file
 
unimernet/datasets/__pycache__/data_utils.cpython-310.pyc ADDED
Binary file (8.18 kB). View file
 
unimernet/datasets/builders/__init__.py ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ from unimernet.datasets.builders.base_dataset_builder import load_dataset_config
9
+ from unimernet.common.registry import registry
10
+ from unimernet.datasets.builders.formula import FormulaRecTrainBuilder, FormulaRecEvalBuilder, \
11
+ MultiScaleFormulaRecTrainBuilder
12
+
13
+ __all__ = [
14
+ "FormulaRecTrainBuilder",
15
+ "FormulaRecEvalBuilder",
16
+ "MultiScaleFormulaRecTrainBuilder",
17
+ ]
18
+
19
+
20
+ def load_dataset(name, cfg_path=None, vis_path=None, data_type=None):
21
+ """
22
+ Example
23
+
24
+ >>> dataset = load_dataset("coco_caption", cfg=None)
25
+ >>> splits = dataset.keys()
26
+ >>> print([len(dataset[split]) for split in splits])
27
+
28
+ """
29
+ if cfg_path is None:
30
+ cfg = None
31
+ else:
32
+ cfg = load_dataset_config(cfg_path)
33
+
34
+ try:
35
+ builder = registry.get_builder_class(name)(cfg)
36
+ except TypeError:
37
+ print(
38
+ f"Dataset {name} not found. Available datasets:\n"
39
+ + ", ".join([str(k) for k in dataset_zoo.get_names()])
40
+ )
41
+ exit(1)
42
+
43
+ if vis_path is not None:
44
+ if data_type is None:
45
+ # use default data type in the config
46
+ data_type = builder.config.data_type
47
+
48
+ assert (
49
+ data_type in builder.config.build_info
50
+ ), f"Invalid data_type {data_type} for {name}."
51
+
52
+ builder.config.build_info.get(data_type).storage = vis_path
53
+
54
+ dataset = builder.build_datasets()
55
+ return dataset
56
+
57
+
58
+ class DatasetZoo:
59
+ def __init__(self) -> None:
60
+ self.dataset_zoo = {
61
+ k: list(v.DATASET_CONFIG_DICT.keys())
62
+ for k, v in sorted(registry.mapping["builder_name_mapping"].items())
63
+ }
64
+
65
+ def get_names(self):
66
+ return list(self.dataset_zoo.keys())
67
+
68
+
69
+ dataset_zoo = DatasetZoo()
unimernet/datasets/builders/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (2.4 kB). View file
 
unimernet/datasets/builders/__pycache__/base_dataset_builder.cpython-310.pyc ADDED
Binary file (6.05 kB). View file
 
unimernet/datasets/builders/__pycache__/formula.cpython-310.pyc ADDED
Binary file (2.44 kB). View file
 
unimernet/datasets/builders/base_dataset_builder.py ADDED
@@ -0,0 +1,233 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ import logging
9
+ import os
10
+ import shutil
11
+ import warnings
12
+
13
+ import unimernet.common.utils as utils
14
+ import torch.distributed as dist
15
+ from unimernet.common.dist_utils import is_dist_avail_and_initialized, is_main_process
16
+ from unimernet.common.registry import registry
17
+ from unimernet.processors.base_processor import BaseProcessor
18
+ from omegaconf import OmegaConf
19
+ from torchvision.datasets.utils import download_url
20
+
21
+
22
+ class BaseDatasetBuilder:
23
+ train_dataset_cls, eval_dataset_cls = None, None
24
+
25
+ def __init__(self, cfg=None):
26
+ super().__init__()
27
+
28
+ if cfg is None:
29
+ # help to create datasets from default config.
30
+ self.config = load_dataset_config(self.default_config_path())
31
+ elif isinstance(cfg, str):
32
+ self.config = load_dataset_config(cfg)
33
+ else:
34
+ # when called from task.build_dataset()
35
+ self.config = cfg
36
+
37
+ self.data_type = self.config.data_type
38
+
39
+ self.vis_processors = {"train": BaseProcessor(), "eval": BaseProcessor()}
40
+ self.text_processors = {"train": BaseProcessor(), "eval": BaseProcessor()}
41
+
42
+ def build_datasets(self):
43
+ # download, split, etc...
44
+ # only called on 1 GPU/TPU in distributed
45
+
46
+ if is_main_process():
47
+ self._download_data()
48
+
49
+ if is_dist_avail_and_initialized():
50
+ dist.barrier()
51
+
52
+ # at this point, all the annotations and image/videos should be all downloaded to the specified locations.
53
+ logging.info("Building datasets...")
54
+ datasets = self.build() # dataset['train'/'val'/'test']
55
+
56
+ return datasets
57
+
58
+ def build_processors(self):
59
+ vis_proc_cfg = self.config.get("vis_processor")
60
+ txt_proc_cfg = self.config.get("text_processor")
61
+
62
+ if vis_proc_cfg is not None:
63
+ vis_train_cfg = vis_proc_cfg.get("train")
64
+ vis_eval_cfg = vis_proc_cfg.get("eval")
65
+
66
+ self.vis_processors["train"] = self._build_proc_from_cfg(vis_train_cfg)
67
+ self.vis_processors["eval"] = self._build_proc_from_cfg(vis_eval_cfg)
68
+
69
+ if txt_proc_cfg is not None:
70
+ txt_train_cfg = txt_proc_cfg.get("train")
71
+ txt_eval_cfg = txt_proc_cfg.get("eval")
72
+
73
+ self.text_processors["train"] = self._build_proc_from_cfg(txt_train_cfg)
74
+ self.text_processors["eval"] = self._build_proc_from_cfg(txt_eval_cfg)
75
+
76
+ @staticmethod
77
+ def _build_proc_from_cfg(cfg):
78
+ return (
79
+ registry.get_processor_class(cfg.name).from_config(cfg)
80
+ if cfg is not None
81
+ else None
82
+ )
83
+
84
+ @classmethod
85
+ def default_config_path(cls, type="default"):
86
+ return utils.get_abs_path(cls.DATASET_CONFIG_DICT[type])
87
+
88
+ def _download_data(self):
89
+ self._download_ann()
90
+ self._download_vis()
91
+
92
+ def _download_ann(self):
93
+ """
94
+ Download annotation files if necessary.
95
+ All the vision-language datasets should have annotations of unified format.
96
+
97
+ storage_path can be:
98
+ (1) relative/absolute: will be prefixed with env.cache_root to make full path if relative.
99
+ (2) basename/dirname: will be suffixed with base name of URL if dirname is provided.
100
+
101
+ Local annotation paths should be relative.
102
+ """
103
+ anns = self.config.build_info.annotations
104
+
105
+ splits = anns.keys()
106
+
107
+ cache_root = registry.get_path("cache_root")
108
+
109
+ for split in splits:
110
+ info = anns[split]
111
+
112
+ urls, storage_paths = info.get("url", None), info.storage
113
+
114
+ if isinstance(urls, str):
115
+ urls = [urls]
116
+ if isinstance(storage_paths, str):
117
+ storage_paths = [storage_paths]
118
+
119
+ assert len(urls) == len(storage_paths)
120
+
121
+ for url_or_filename, storage_path in zip(urls, storage_paths):
122
+ # if storage_path is relative, make it full by prefixing with cache_root.
123
+ if not os.path.isabs(storage_path):
124
+ storage_path = os.path.join(cache_root, storage_path)
125
+
126
+ dirname = os.path.dirname(storage_path)
127
+ if not os.path.exists(dirname):
128
+ os.makedirs(dirname)
129
+
130
+ if os.path.isfile(url_or_filename):
131
+ src, dst = url_or_filename, storage_path
132
+ if not os.path.exists(dst):
133
+ shutil.copyfile(src=src, dst=dst)
134
+ else:
135
+ logging.info("Using existing file {}.".format(dst))
136
+ else:
137
+ if os.path.isdir(storage_path):
138
+ # if only dirname is provided, suffix with basename of URL.
139
+ raise ValueError(
140
+ "Expecting storage_path to be a file path, got directory {}".format(
141
+ storage_path
142
+ )
143
+ )
144
+ else:
145
+ filename = os.path.basename(storage_path)
146
+
147
+ download_url(url=url_or_filename, root=dirname, filename=filename)
148
+
149
+ def _download_vis(self):
150
+
151
+ storage_path = self.config.build_info.get(self.data_type).storage
152
+ storage_path = utils.get_cache_path(storage_path)
153
+
154
+ if not os.path.exists(storage_path):
155
+ warnings.warn(
156
+ f"""
157
+ The specified path {storage_path} for visual inputs does not exist.
158
+ Please provide a correct path to the visual inputs or
159
+ refer to datasets/download_scripts/README.md for downloading instructions.
160
+ """
161
+ )
162
+
163
+ def build(self):
164
+ """
165
+ Create by split datasets inheriting torch.utils.data.Datasets.
166
+
167
+ # build() can be dataset-specific. Overwrite to customize.
168
+ """
169
+ self.build_processors()
170
+
171
+ build_info = self.config.build_info
172
+
173
+ ann_info = build_info.annotations
174
+ vis_info = build_info.get(self.data_type)
175
+
176
+ datasets = dict()
177
+ for split in ann_info.keys():
178
+ if split not in ["train", "val", "test"]:
179
+ continue
180
+
181
+ is_train = split == "train"
182
+
183
+ # processors
184
+ vis_processor = (
185
+ self.vis_processors["train"]
186
+ if is_train
187
+ else self.vis_processors["eval"]
188
+ )
189
+ text_processor = (
190
+ self.text_processors["train"]
191
+ if is_train
192
+ else self.text_processors["eval"]
193
+ )
194
+
195
+ # annotation path
196
+ ann_paths = ann_info.get(split).storage
197
+ if isinstance(ann_paths, str):
198
+ ann_paths = [ann_paths]
199
+
200
+ abs_ann_paths = []
201
+ for ann_path in ann_paths:
202
+ if not os.path.isabs(ann_path):
203
+ ann_path = utils.get_cache_path(ann_path)
204
+ abs_ann_paths.append(ann_path)
205
+ ann_paths = abs_ann_paths
206
+
207
+ # visual data storage path
208
+ vis_path = vis_info.storage
209
+
210
+ if not os.path.isabs(vis_path):
211
+ # vis_path = os.path.join(utils.get_cache_path(), vis_path)
212
+ vis_path = utils.get_cache_path(vis_path)
213
+
214
+ if not os.path.exists(vis_path):
215
+ warnings.warn("storage path {} does not exist.".format(vis_path))
216
+
217
+ # create datasets
218
+ dataset_cls = self.train_dataset_cls if is_train else self.eval_dataset_cls
219
+ datasets[split] = dataset_cls(
220
+ vis_processor=vis_processor,
221
+ text_processor=text_processor,
222
+ ann_paths=ann_paths,
223
+ vis_root=vis_path,
224
+ )
225
+
226
+ return datasets
227
+
228
+
229
+ def load_dataset_config(cfg_path):
230
+ cfg = OmegaConf.load(cfg_path).datasets
231
+ cfg = cfg[list(cfg.keys())[0]]
232
+
233
+ return cfg
unimernet/datasets/builders/formula.py ADDED
@@ -0,0 +1,105 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import logging
2
+ from unimernet.common.registry import registry
3
+ from unimernet.datasets.builders.base_dataset_builder import BaseDatasetBuilder
4
+ from unimernet.datasets.datasets.formula import Im2LatexDataset
5
+ from unimernet.datasets.datasets.formula_multi_scale import MultiScaleIm2LatexDataset
6
+
7
+
8
+ @registry.register_builder("formula_rec_train")
9
+ class FormulaRecTrainBuilder(BaseDatasetBuilder):
10
+ train_dataset_cls = Im2LatexDataset
11
+ DATASET_CONFIG_DICT = {
12
+ "default": "configs/datasets/formula/formula_train.yaml"
13
+ }
14
+ LOG_INFO = "Formula Recgnition Train"
15
+
16
+ def build_datasets(self):
17
+ logging.info(f"Building {self.LOG_INFO} datasets ...")
18
+ self.build_processors()
19
+
20
+ build_info = self.config.build_info
21
+ anno_path = build_info.annotation,
22
+ vis_root = build_info.images
23
+ anno_path = [anno_path] if isinstance(anno_path, str) else anno_path
24
+ vis_root = [vis_root] if isinstance(vis_root, str) else vis_root
25
+ datasets = dict()
26
+
27
+ # create datasets
28
+ dataset_cls = self.train_dataset_cls
29
+ datasets['train'] = dataset_cls(
30
+ vis_processor=self.vis_processors["train"],
31
+ text_processor=self.text_processors["train"],
32
+ vis_root=vis_root,
33
+ anno_path=anno_path,
34
+ )
35
+ print(datasets['train'][0])
36
+
37
+ return datasets
38
+
39
+
40
+ @registry.register_builder("multi_scale_formula_rec_train")
41
+ class MultiScaleFormulaRecTrainBuilder(BaseDatasetBuilder):
42
+ train_dataset_cls = MultiScaleIm2LatexDataset
43
+ DATASET_CONFIG_DICT = {
44
+ "default": "configs/datasets/formula/multi_scale_formula_train.yaml"
45
+ }
46
+ LOG_INFO = "Multi Scale Formula Recgnition Train"
47
+
48
+ def build_datasets(self):
49
+ logging.info(f"Building {self.LOG_INFO} datasets ...")
50
+ self.build_processors()
51
+
52
+ build_info = self.config.build_info
53
+ anno_path = build_info.annotation,
54
+ vis_root = build_info.images
55
+
56
+ anno_path = [anno_path] if isinstance(anno_path, str) else anno_path
57
+ vis_root = [vis_root] if isinstance(vis_root, str) else vis_root
58
+
59
+ datasets = dict()
60
+
61
+ # create datasets
62
+ dataset_cls = self.train_dataset_cls
63
+ datasets['train'] = dataset_cls(
64
+ vis_processor=self.vis_processors["train"],
65
+ text_processor=self.text_processors["train"],
66
+ vis_root=vis_root,
67
+ anno_path=anno_path,
68
+ )
69
+ print(datasets['train'][0])
70
+
71
+ return datasets
72
+
73
+
74
+ @registry.register_builder("formula_rec_eval")
75
+ class FormulaRecEvalBuilder(BaseDatasetBuilder):
76
+ eval_dataset_cls = Im2LatexDataset
77
+ DATASET_CONFIG_DICT = {
78
+ "default": "configs/datasets/formula/formula_eval.yaml"
79
+ }
80
+ LOG_INFO = "Formula Recgnition Eval"
81
+
82
+ def build_datasets(self):
83
+ logging.info(f"Building {self.LOG_INFO} datasets ...")
84
+ self.build_processors()
85
+
86
+ build_info = self.config.build_info
87
+ anno_path = build_info.annotation,
88
+ vis_root = build_info.images
89
+
90
+ anno_path = [anno_path] if isinstance(anno_path, str) else anno_path
91
+ vis_root = [vis_root] if isinstance(vis_root, str) else vis_root
92
+
93
+ datasets = dict()
94
+
95
+ # create datasets
96
+ dataset_cls = self.eval_dataset_cls
97
+ datasets['eval'] = dataset_cls(
98
+ vis_processor=self.vis_processors["eval"],
99
+ text_processor=self.text_processors["eval"],
100
+ vis_root=vis_root,
101
+ anno_path=anno_path,
102
+ )
103
+ print(datasets['eval'][0])
104
+
105
+ return datasets
unimernet/datasets/data_utils.py ADDED
@@ -0,0 +1,284 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ import gzip
9
+ import logging
10
+ import os
11
+ import random as rnd
12
+ import tarfile
13
+ import zipfile
14
+
15
+ import decord
16
+ import webdataset as wds
17
+ import numpy as np
18
+ import torch
19
+ from torch.utils.data.dataset import IterableDataset, ChainDataset
20
+ from decord import VideoReader
21
+ from unimernet.common.registry import registry
22
+ from unimernet.datasets.datasets.base_dataset import ConcatDataset
23
+ from tqdm import tqdm
24
+
25
+ decord.bridge.set_bridge("torch")
26
+ MAX_INT = registry.get("MAX_INT")
27
+
28
+
29
+ def load_video(video_path, n_frms=MAX_INT, height=-1, width=-1, sampling="uniform"):
30
+ vr = VideoReader(uri=video_path, height=height, width=width)
31
+
32
+ vlen = len(vr)
33
+ start, end = 0, vlen
34
+
35
+ n_frms = min(n_frms, vlen)
36
+
37
+ if sampling == "uniform":
38
+ indices = np.arange(start, end, vlen / n_frms).astype(int)
39
+ elif sampling == "headtail":
40
+ indices_h = sorted(rnd.sample(range(vlen // 2), n_frms // 2))
41
+ indices_t = sorted(rnd.sample(range(vlen // 2, vlen), n_frms // 2))
42
+ indices = indices_h + indices_t
43
+ else:
44
+ raise NotImplementedError
45
+
46
+ # get_batch -> T, H, W, C
47
+ frms = vr.get_batch(indices).permute(3, 0, 1, 2).float() # (C, T, H, W)
48
+
49
+ return frms
50
+
51
+
52
+ def apply_to_sample(f, sample):
53
+ if len(sample) == 0:
54
+ return {}
55
+
56
+ def _apply(x):
57
+ if torch.is_tensor(x):
58
+ return f(x)
59
+ elif isinstance(x, dict):
60
+ return {key: _apply(value) for key, value in x.items()}
61
+ elif isinstance(x, list):
62
+ return [_apply(x) for x in x]
63
+ else:
64
+ return x
65
+
66
+ return _apply(sample)
67
+
68
+
69
+ def move_to_cuda(sample):
70
+ def _move_to_cuda(tensor):
71
+ return tensor.cuda()
72
+
73
+ return apply_to_sample(_move_to_cuda, sample)
74
+
75
+
76
+ def prepare_sample(samples, cuda_enabled=True):
77
+ if cuda_enabled:
78
+ samples = move_to_cuda(samples)
79
+
80
+ # TODO fp16 support
81
+
82
+ return samples
83
+
84
+
85
+ def reorg_datasets_by_split(datasets):
86
+ """
87
+ Organizes datasets by split.
88
+
89
+ Args:
90
+ datasets: dict of torch.utils.data.Dataset objects by name.
91
+
92
+ Returns:
93
+ Dict of datasets by split {split_name: List[Datasets]}.
94
+ """
95
+ # if len(datasets) == 1:
96
+ # return datasets[list(datasets.keys())[0]]
97
+ # else:
98
+ reorg_datasets = dict()
99
+
100
+ # reorganize by split
101
+ for _, dataset in datasets.items():
102
+ for split_name, dataset_split in dataset.items():
103
+ if split_name not in reorg_datasets:
104
+ reorg_datasets[split_name] = [dataset_split]
105
+ else:
106
+ reorg_datasets[split_name].append(dataset_split)
107
+
108
+ return reorg_datasets
109
+
110
+
111
+ def concat_datasets(datasets):
112
+ """
113
+ Concatenates multiple datasets into a single dataset.
114
+
115
+ It supports may-style datasets and DataPipeline from WebDataset. Currently, does not support
116
+ generic IterableDataset because it requires creating separate samplers.
117
+
118
+ Now only supports conctenating training datasets and assuming validation and testing
119
+ have only a single dataset. This is because metrics should not be computed on the concatenated
120
+ datasets.
121
+
122
+ Args:
123
+ datasets: dict of torch.utils.data.Dataset objects by split.
124
+
125
+ Returns:
126
+ Dict of concatenated datasets by split, "train" is the concatenation of multiple datasets,
127
+ "val" and "test" remain the same.
128
+
129
+ If the input training datasets contain both map-style and DataPipeline datasets, returns
130
+ a tuple, where the first element is a concatenated map-style dataset and the second
131
+ element is a chained DataPipeline dataset.
132
+
133
+ """
134
+ # concatenate datasets in the same split
135
+ for split_name in datasets:
136
+ if split_name != "train":
137
+ assert (
138
+ len(datasets[split_name]) == 1
139
+ ), "Do not support multiple {} datasets.".format(split_name)
140
+ datasets[split_name] = datasets[split_name][0]
141
+ else:
142
+ iterable_datasets, map_datasets = [], []
143
+ for dataset in datasets[split_name]:
144
+ if isinstance(dataset, wds.DataPipeline):
145
+ logging.info(
146
+ "Dataset {} is IterableDataset, can't be concatenated.".format(
147
+ dataset
148
+ )
149
+ )
150
+ iterable_datasets.append(dataset)
151
+ elif isinstance(dataset, IterableDataset):
152
+ raise NotImplementedError(
153
+ "Do not support concatenation of generic IterableDataset."
154
+ )
155
+ else:
156
+ map_datasets.append(dataset)
157
+
158
+ # if len(iterable_datasets) > 0:
159
+ # concatenate map-style datasets and iterable-style datasets separately
160
+ chained_datasets = (
161
+ ChainDataset(iterable_datasets) if len(iterable_datasets) > 0 else None
162
+ )
163
+ concat_datasets = (
164
+ ConcatDataset(map_datasets) if len(map_datasets) > 0 else None
165
+ )
166
+
167
+ train_datasets = concat_datasets, chained_datasets
168
+ train_datasets = tuple([x for x in train_datasets if x is not None])
169
+ train_datasets = (
170
+ train_datasets[0] if len(train_datasets) == 1 else train_datasets
171
+ )
172
+
173
+ datasets[split_name] = train_datasets
174
+
175
+ return datasets
176
+
177
+
178
+ def extract_archive(from_path, to_path=None, overwrite=False):
179
+ """Extract archive.
180
+
181
+ Args:
182
+ from_path: the path of the archive.
183
+ to_path: the root path of the extracted files (directory of from_path)
184
+ overwrite: overwrite existing files (False)
185
+
186
+ Returns:
187
+ List of paths to extracted files even if not overwritten.
188
+
189
+ Examples:
190
+ >>> url = 'http://www.quest.dcs.shef.ac.uk/wmt16_files_mmt/validation.tar.gz'
191
+ >>> from_path = './validation.tar.gz'
192
+ >>> to_path = './'
193
+ >>> torchtext.utils.download_from_url(url, from_path)
194
+ >>> torchtext.utils.extract_archive(from_path, to_path)
195
+ >>> ['.data/val.de', '.data/val.en']
196
+ >>> torchtext.utils.download_from_url(url, from_path)
197
+ >>> torchtext.utils.extract_archive(from_path, to_path)
198
+ >>> ['.data/val.de', '.data/val.en']
199
+
200
+ """
201
+
202
+ if to_path is None:
203
+ to_path = os.path.dirname(from_path)
204
+
205
+ if from_path.endswith((".tar.gz", ".tgz")):
206
+ logging.info("Opening tar file {} to {}.".format(from_path, to_path))
207
+ with tarfile.open(from_path, "r") as tar:
208
+ files = []
209
+ for file_ in tqdm(tar):
210
+ file_path = os.path.join(to_path, file_.name)
211
+ if file_.isfile():
212
+ files.append(file_path)
213
+ if os.path.exists(file_path):
214
+ logging.info("{} already extracted.".format(file_path))
215
+ if not overwrite:
216
+ continue
217
+ tar.extract(file_, to_path)
218
+ logging.info("Finished extracting tar file {}.".format(from_path))
219
+ return files
220
+
221
+ elif from_path.endswith(".zip"):
222
+ assert zipfile.is_zipfile(from_path), from_path
223
+ logging.info("Opening zip file {} to {}.".format(from_path, to_path))
224
+ with zipfile.ZipFile(from_path, "r") as zfile:
225
+ files = []
226
+ for file_ in tqdm(zfile.namelist()):
227
+ file_path = os.path.join(to_path, file_)
228
+ files.append(file_path)
229
+ if os.path.exists(file_path):
230
+ logging.info("{} already extracted.".format(file_path))
231
+ if not overwrite:
232
+ continue
233
+ zfile.extract(file_, to_path)
234
+ files = [f for f in files if os.path.isfile(f)]
235
+ logging.info("Finished extracting zip file {}.".format(from_path))
236
+ return files
237
+
238
+ elif from_path.endswith(".gz"):
239
+ logging.info("Opening gz file {} to {}.".format(from_path, to_path))
240
+ default_block_size = 65536
241
+ filename = from_path[:-3]
242
+ files = [filename]
243
+ with gzip.open(from_path, "rb") as gzfile, open(filename, "wb") as d_file:
244
+ while True:
245
+ block = gzfile.read(default_block_size)
246
+ if not block:
247
+ break
248
+ else:
249
+ d_file.write(block)
250
+ d_file.write(block)
251
+ logging.info("Finished extracting gz file {}.".format(from_path))
252
+ return files
253
+
254
+ else:
255
+ raise NotImplementedError(
256
+ "We currently only support tar.gz, .tgz, .gz and zip achives."
257
+ )
258
+
259
+
260
+ def save_frames_grid(img_array, out_path):
261
+ import torch
262
+ from PIL import Image
263
+ from torchvision.utils import make_grid
264
+
265
+ if len(img_array.shape) == 3:
266
+ img_array = img_array.unsqueeze(0)
267
+ elif len(img_array.shape) == 5:
268
+ b, t, c, h, w = img_array.shape
269
+ img_array = img_array.view(-1, c, h, w)
270
+ elif len(img_array.shape) == 4:
271
+ pass
272
+ else:
273
+ raise NotImplementedError(
274
+ "Supports only (b,t,c,h,w)-shaped inputs. First two dimensions can be ignored."
275
+ )
276
+
277
+ assert img_array.shape[1] == 3, "Exepcting input shape of (H, W, 3), i.e. RGB-only."
278
+
279
+ grid = make_grid(img_array)
280
+ ndarr = grid.permute(1, 2, 0).to("cpu", torch.uint8).numpy()
281
+
282
+ img = Image.fromarray(ndarr)
283
+
284
+ img.save(out_path)
unimernet/datasets/datasets/__pycache__/base_dataset.cpython-310.pyc ADDED
Binary file (3.55 kB). View file
 
unimernet/datasets/datasets/__pycache__/formula.cpython-310.pyc ADDED
Binary file (2.85 kB). View file
 
unimernet/datasets/datasets/__pycache__/formula_multi_scale.cpython-310.pyc ADDED
Binary file (1.23 kB). View file
 
unimernet/datasets/datasets/base_dataset.py ADDED
@@ -0,0 +1,103 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ from PIL import Image, ImageFile
3
+ import os.path as osp
4
+
5
+ ImageFile.LOAD_TRUNCATED_IMAGES = True
6
+
7
+ from io import BytesIO
8
+ from typing import Iterable
9
+ from torch.utils.data import Dataset, ConcatDataset
10
+ import torch
11
+
12
+
13
+ class BaseDataset(Dataset):
14
+
15
+ def __init__(self, vis_processor, text_processor, vis_root, anno_path):
16
+
17
+ self.vis_root = vis_root
18
+ # if isinstance(anno_path, tuple) or isinstance(anno_path, list):
19
+ # anno_path = anno_path[0]
20
+ self.anno_path = anno_path
21
+
22
+ self.vis_processor = vis_processor
23
+ self.text_processor = text_processor
24
+
25
+ self.samples = self.init_samples()
26
+ self.reader = self.init_reader()
27
+
28
+ print('total {} {} samples'.format(self.__len__(), self.__class__.__name__))
29
+
30
+ for idx in range(10):
31
+ self.__getitem__(idx)
32
+
33
+ def __len__(self):
34
+ return len(self.samples)
35
+
36
+ def __getitem__(self, index):
37
+ raise NotImplementedError
38
+
39
+ def init_samples(self):
40
+ # read annotation from ceph
41
+ if self.anno_path.startswith('cluster'):
42
+ from petrel_client.client import Client
43
+ client = Client("~/petreloss.conf")
44
+ samples = json.loads(client.get(self.anno_path))
45
+ else:
46
+ samples = json.load(open(self.anno_path, 'r'))
47
+ return samples
48
+
49
+ def init_reader(self):
50
+ if self.vis_root.startswith('cluster'):
51
+ from petrel_client.client import Client
52
+ client = Client("~/petreloss.conf")
53
+ reader = {'type': 'PetrelReader', 'body': client.get}
54
+ else:
55
+ reader = {'type': 'LocalReader', 'body': Image.open}
56
+ return reader
57
+
58
+ def _read_image(self, sample, image_key="image"):
59
+ img_file = sample[image_key]
60
+ image_path = osp.join(self.vis_root, img_file)
61
+ image = self.reader['body'](image_path)
62
+ if isinstance(image, bytes):
63
+ bytes_stream = BytesIO(image)
64
+ image = Image.open(bytes_stream)
65
+ image = image.convert("RGB")
66
+ return image
67
+
68
+ def collater(self, samples):
69
+ image_list, question_list, answer_list = [], [], []
70
+
71
+ for sample in samples:
72
+ image_list.append(sample["image"])
73
+ question_list.append(sample["text_input"])
74
+ answer_list.append(sample["text_output"])
75
+
76
+ return {
77
+ "image": torch.stack(image_list, dim=0),
78
+ "text_input": question_list,
79
+ "text_output": answer_list,
80
+ "data_type": "vqa",
81
+ }
82
+
83
+
84
+ class ConcatDataset(ConcatDataset):
85
+ def __init__(self, datasets: Iterable[Dataset]) -> None:
86
+ super().__init__(datasets)
87
+
88
+ def collater(self, samples):
89
+ # TODO For now only supports datasets with same underlying collater implementations
90
+
91
+ all_keys = set()
92
+ for s in samples:
93
+ all_keys.update(s)
94
+
95
+ shared_keys = all_keys
96
+ for s in samples:
97
+ shared_keys = shared_keys & set(s.keys())
98
+
99
+ samples_shared_keys = []
100
+ for s in samples:
101
+ samples_shared_keys.append({k: s[k] for k in s.keys() if k in shared_keys})
102
+
103
+ return self.datasets[0].collater(samples_shared_keys)
unimernet/datasets/datasets/dataloader_utils.py ADDED
@@ -0,0 +1,200 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ import time
9
+ import random
10
+ import torch
11
+ from unimernet.datasets.data_utils import move_to_cuda
12
+ from torch.utils.data import DataLoader
13
+
14
+
15
+ class MultiIterLoader:
16
+ """
17
+ A simple wrapper for iterating over multiple iterators.
18
+
19
+ Args:
20
+ loaders (List[Loader]): List of Iterator loaders.
21
+ ratios (List[float]): List of ratios to sample from each loader. If None, all loaders are sampled uniformly.
22
+ """
23
+
24
+ def __init__(self, loaders, ratios=None):
25
+ # assert all loaders has __next__ method
26
+ for loader in loaders:
27
+ assert hasattr(
28
+ loader, "__next__"
29
+ ), "Loader {} has no __next__ method.".format(loader)
30
+
31
+ if ratios is None:
32
+ ratios = [1.0] * len(loaders)
33
+ else:
34
+ assert len(ratios) == len(loaders)
35
+ ratios = [float(ratio) / sum(ratios) for ratio in ratios]
36
+
37
+ self.loaders = loaders
38
+ self.ratios = ratios
39
+
40
+ def __next__(self):
41
+ # random sample from each loader by ratio
42
+ loader_idx = random.choices(range(len(self.loaders)), self.ratios, k=1)[0]
43
+ return next(self.loaders[loader_idx])
44
+
45
+ def __len__(self):
46
+ return sum([len(_) for _ in self.loaders if hasattr(_, "__len__")])
47
+
48
+
49
+ class ConcatLoader:
50
+ """
51
+ A simple wrapper for iterating over multiple iterators.
52
+
53
+ Args:
54
+ loaders (List[Loader]): List of Iterator loaders.
55
+ ratios (List[float]): List of ratios to sample from each loader. If None, all loaders are sampled uniformly.
56
+ """
57
+
58
+ def __init__(self, loaders):
59
+ # assert all loaders has __next__ method
60
+ for loader in loaders:
61
+ assert hasattr(
62
+ loader, "__len__"
63
+ ), "Loader {} has no __len__ method.".format(loader)
64
+
65
+ self._epoch = 0
66
+ self._loader_lens = [len(_) for _ in loaders]
67
+ self._rest_lens = self._loader_lens.copy()
68
+
69
+ self.loaders = loaders
70
+
71
+ def __next__(self):
72
+ # random sample from each loader by ratio
73
+ loader_idx = random.choices(range(len(self.loaders)), self._rest_lens, k=1)[0]
74
+ self._rest_lens[loader_idx] -= 1
75
+ if sum(self._rest_lens) == 0:
76
+ self._epoch += 1
77
+ self._rest_lens = self._loader_lens.copy()
78
+ return next(self.loaders[loader_idx])
79
+
80
+ def __len__(self):
81
+ return sum([len(_) for _ in self.loaders if hasattr(_, "__len__")])
82
+
83
+
84
+ class PrefetchLoader(object):
85
+ """
86
+ Modified from https://github.com/ChenRocks/UNITER.
87
+
88
+ overlap compute and cuda data transfer
89
+ (copied and then modified from nvidia apex)
90
+ """
91
+
92
+ def __init__(self, loader):
93
+ self.loader = loader
94
+ self.stream = torch.cuda.Stream()
95
+
96
+ def __iter__(self):
97
+ loader_it = iter(self.loader)
98
+ self.preload(loader_it)
99
+ batch = self.next(loader_it)
100
+ while batch is not None:
101
+ is_tuple = isinstance(batch, tuple)
102
+ if is_tuple:
103
+ task, batch = batch
104
+
105
+ if is_tuple:
106
+ yield task, batch
107
+ else:
108
+ yield batch
109
+ batch = self.next(loader_it)
110
+
111
+ def __len__(self):
112
+ return len(self.loader)
113
+
114
+ def preload(self, it):
115
+ try:
116
+ self.batch = next(it)
117
+ except StopIteration:
118
+ self.batch = None
119
+ return
120
+ # if record_stream() doesn't work, another option is to make sure
121
+ # device inputs are created on the main stream.
122
+ # self.next_input_gpu = torch.empty_like(self.next_input,
123
+ # device='cuda')
124
+ # self.next_target_gpu = torch.empty_like(self.next_target,
125
+ # device='cuda')
126
+ # Need to make sure the memory allocated for next_* is not still in use
127
+ # by the main stream at the time we start copying to next_*:
128
+ # self.stream.wait_stream(torch.cuda.current_stream())
129
+ with torch.cuda.stream(self.stream):
130
+ self.batch = move_to_cuda(self.batch)
131
+ # more code for the alternative if record_stream() doesn't work:
132
+ # copy_ will record the use of the pinned source tensor in this
133
+ # side stream.
134
+ # self.next_input_gpu.copy_(self.next_input, non_blocking=True)
135
+ # self.next_target_gpu.copy_(self.next_target, non_blocking=True)
136
+ # self.next_input = self.next_input_gpu
137
+ # self.next_target = self.next_target_gpu
138
+
139
+ def next(self, it):
140
+ torch.cuda.current_stream().wait_stream(self.stream)
141
+ batch = self.batch
142
+ if batch is not None:
143
+ record_cuda_stream(batch)
144
+ self.preload(it)
145
+ return batch
146
+
147
+ def __getattr__(self, name):
148
+ method = self.loader.__getattribute__(name)
149
+ return method
150
+
151
+
152
+ def record_cuda_stream(batch):
153
+ if isinstance(batch, torch.Tensor):
154
+ batch.record_stream(torch.cuda.current_stream())
155
+ elif isinstance(batch, list) or isinstance(batch, tuple):
156
+ for t in batch:
157
+ record_cuda_stream(t)
158
+ elif isinstance(batch, dict):
159
+ for t in batch.values():
160
+ record_cuda_stream(t)
161
+ else:
162
+ pass
163
+
164
+
165
+ class IterLoader:
166
+ """
167
+ A wrapper to convert DataLoader as an infinite iterator.
168
+
169
+ Modified from:
170
+ https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/iter_based_runner.py
171
+ """
172
+
173
+ def __init__(self, dataloader: DataLoader, use_distributed: bool = False):
174
+ self._dataloader = dataloader
175
+ self.iter_loader = iter(self._dataloader)
176
+ self._use_distributed = use_distributed
177
+ self._epoch = 0
178
+
179
+ @property
180
+ def epoch(self) -> int:
181
+ return self._epoch
182
+
183
+ def __next__(self):
184
+ try:
185
+ data = next(self.iter_loader)
186
+ except StopIteration:
187
+ self._epoch += 1
188
+ if hasattr(self._dataloader.sampler, "set_epoch") and self._use_distributed:
189
+ self._dataloader.sampler.set_epoch(self._epoch)
190
+ time.sleep(2) # Prevent possible deadlock during epoch transition
191
+ self.iter_loader = iter(self._dataloader)
192
+ data = next(self.iter_loader)
193
+
194
+ return data
195
+
196
+ def __iter__(self):
197
+ return self
198
+
199
+ def __len__(self):
200
+ return len(self._dataloader)
unimernet/datasets/datasets/formula.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from .base_dataset import BaseDataset
3
+ import os.path as osp
4
+ import glob
5
+ from io import BytesIO
6
+ from PIL import Image
7
+
8
+
9
+ class Im2LatexDataset(BaseDataset):
10
+
11
+ def init_samples(self):
12
+ samples = []
13
+ for vis_root, anno_path in zip(self.vis_root, self.anno_path):
14
+ images = [path.replace('\\', '/') for path in glob.glob(osp.join(vis_root, '*.png'))]
15
+ indices = [int(osp.basename(img).split('.')[0]) for img in images]
16
+
17
+ eqs = open(anno_path, 'r').read().split('\n')
18
+ eqs = [eqs[_] for _ in indices]
19
+
20
+ for i, e in zip(images, eqs):
21
+ samples.append({"image": i, "equation": e, "vis_root": vis_root})
22
+ return samples
23
+
24
+ def __getitem__(self, index):
25
+ ann = self.samples[index]
26
+ try:
27
+ image = self.vis_processor(self._read_image(ann))
28
+ except Exception:
29
+ return self[(index + 1) % len(self)]
30
+ if image is None:
31
+ return self[(index + 1) % len(self)]
32
+ equation = ann["equation"]
33
+ return {"image": image, "text_input": equation, "id": index}
34
+
35
+ def _read_image(self, sample, image_key="image"):
36
+ img_file = sample[image_key]
37
+ vis_root = sample["vis_root"]
38
+ image_path = osp.join(vis_root, img_file)
39
+ image = self.reader['body'](image_path)
40
+ if isinstance(image, bytes):
41
+ bytes_stream = BytesIO(image)
42
+ image = Image.open(bytes_stream)
43
+ image = image.convert("RGB")
44
+ return image
45
+
46
+ def init_reader(self):
47
+ if not isinstance(self.vis_root, str):
48
+ vis_root = self.vis_root[0]
49
+ else:
50
+ vis_root = self.vis_root
51
+ if vis_root.startswith('cluster'):
52
+ from petrel_client.client import Client
53
+ client = Client("~/petreloss.conf")
54
+ reader = {'type': 'PetrelReader', 'body': client.get}
55
+ else:
56
+ reader = {'type': 'LocalReader', 'body': Image.open}
57
+ return reader
58
+
59
+ def collater(self, samples):
60
+ image_list, question_list, id_list = [], [], []
61
+
62
+ for sample in samples:
63
+ image_list.append(sample["image"])
64
+ question_list.append(sample["text_input"])
65
+ id_list.append(sample["id"])
66
+
67
+ return {
68
+ "image": torch.stack(image_list, dim=0),
69
+ "text_input": question_list,
70
+ "id": id_list
71
+ }
unimernet/datasets/datasets/formula_multi_scale.py ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from .formula import Im2LatexDataset
3
+
4
+
5
+ class MultiScaleIm2LatexDataset(Im2LatexDataset):
6
+
7
+ def __getitem__(self, index):
8
+ ann = self.samples[index]
9
+ try:
10
+ pil_image = self._read_image(ann)
11
+ image = self.vis_processor(pil_image)
12
+ except Exception:
13
+ return self[(index + 1) % len(self)]
14
+ if image is None:
15
+ return self[(index + 1) % len(self)]
16
+ equation = ann["equation"]
17
+ return {"image": image, "text_input": equation, "id": index, "raw_image": pil_image}
18
+
19
+ def collater(self, samples):
20
+ self.vis_processor.reset_scale()
21
+ image_list, question_list, id_list = [], [], []
22
+
23
+ for sample in samples:
24
+ image_list.append(self.vis_processor(sample["raw_image"]))
25
+ question_list.append(sample["text_input"])
26
+ id_list.append(sample["id"])
27
+
28
+ return {
29
+ "image": torch.stack(image_list, dim=0),
30
+ "text_input": question_list,
31
+ "id": id_list
32
+ }
unimernet/models/__init__.py ADDED
@@ -0,0 +1,198 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ import logging
9
+ import torch
10
+ from omegaconf import OmegaConf
11
+ from unimernet.common.registry import registry
12
+
13
+ from unimernet.models.base_model import BaseModel
14
+
15
+ from unimernet.processors.base_processor import BaseProcessor
16
+ from unimernet.models.unimernet.unimernet import UniMERModel
17
+
18
+ __all__ = [
19
+ "load_model",
20
+ "BaseModel",
21
+ "UniMERModel",
22
+ ]
23
+
24
+
25
+ def load_model(name, model_type, is_eval=False, device="cpu", checkpoint=None):
26
+ """
27
+ Load supported models.
28
+
29
+ To list all available models and types in registry:
30
+ >>> from unimernet.models import model_zoo
31
+ >>> print(model_zoo)
32
+
33
+ Args:
34
+ name (str): name of the model.
35
+ model_type (str): type of the model.
36
+ is_eval (bool): whether the model is in eval mode. Default: False.
37
+ device (str): device to use. Default: "cpu".
38
+ checkpoint (str): path or to checkpoint. Default: None.
39
+ Note that expecting the checkpoint to have the same keys in state_dict as the model.
40
+
41
+ Returns:
42
+ model (torch.nn.Module): model.
43
+ """
44
+
45
+ model = registry.get_model_class(name).from_pretrained(model_type=model_type)
46
+
47
+ if checkpoint is not None:
48
+ model.load_checkpoint(checkpoint)
49
+
50
+ if is_eval:
51
+ model.eval()
52
+
53
+ if device == "cpu":
54
+ model = model.float()
55
+
56
+ return model.to(device)
57
+
58
+
59
+ def load_preprocess(config):
60
+ """
61
+ Load preprocessor configs and construct preprocessors.
62
+
63
+ If no preprocessor is specified, return BaseProcessor, which does not do any preprocessing.
64
+
65
+ Args:
66
+ config (dict): preprocessor configs.
67
+
68
+ Returns:
69
+ vis_processors (dict): preprocessors for visual inputs.
70
+ txt_processors (dict): preprocessors for text inputs.
71
+
72
+ Key is "train" or "eval" for processors used in training and evaluation respectively.
73
+ """
74
+
75
+ def _build_proc_from_cfg(cfg):
76
+ return (
77
+ registry.get_processor_class(cfg.name).from_config(cfg)
78
+ if cfg is not None
79
+ else BaseProcessor()
80
+ )
81
+
82
+ vis_processors = dict()
83
+ txt_processors = dict()
84
+
85
+ vis_proc_cfg = config.get("vis_processor")
86
+ txt_proc_cfg = config.get("text_processor")
87
+
88
+ if vis_proc_cfg is not None:
89
+ vis_train_cfg = vis_proc_cfg.get("train")
90
+ vis_eval_cfg = vis_proc_cfg.get("eval")
91
+ else:
92
+ vis_train_cfg = None
93
+ vis_eval_cfg = None
94
+
95
+ vis_processors["train"] = _build_proc_from_cfg(vis_train_cfg)
96
+ vis_processors["eval"] = _build_proc_from_cfg(vis_eval_cfg)
97
+
98
+ if txt_proc_cfg is not None:
99
+ txt_train_cfg = txt_proc_cfg.get("train")
100
+ txt_eval_cfg = txt_proc_cfg.get("eval")
101
+ else:
102
+ txt_train_cfg = None
103
+ txt_eval_cfg = None
104
+
105
+ txt_processors["train"] = _build_proc_from_cfg(txt_train_cfg)
106
+ txt_processors["eval"] = _build_proc_from_cfg(txt_eval_cfg)
107
+
108
+ return vis_processors, txt_processors
109
+
110
+
111
+ def load_model_and_preprocess(name, model_type, is_eval=False, device="cpu"):
112
+ """
113
+ Load model and its related preprocessors.
114
+
115
+ List all available models and types in registry:
116
+ >>> from unimernet.models import model_zoo
117
+ >>> print(model_zoo)
118
+
119
+ Args:
120
+ name (str): name of the model.
121
+ model_type (str): type of the model.
122
+ is_eval (bool): whether the model is in eval mode. Default: False.
123
+ device (str): device to use. Default: "cpu".
124
+
125
+ Returns:
126
+ model (torch.nn.Module): model.
127
+ vis_processors (dict): preprocessors for visual inputs.
128
+ txt_processors (dict): preprocessors for text inputs.
129
+ """
130
+ model_cls = registry.get_model_class(name)
131
+
132
+ # load model
133
+ model = model_cls.from_pretrained(model_type=model_type)
134
+
135
+ if is_eval:
136
+ model.eval()
137
+
138
+ # load preprocess
139
+ cfg = OmegaConf.load(model_cls.default_config_path(model_type))
140
+ if cfg is not None:
141
+ preprocess_cfg = cfg.preprocess
142
+
143
+ vis_processors, txt_processors = load_preprocess(preprocess_cfg)
144
+ else:
145
+ vis_processors, txt_processors = None, None
146
+ logging.info(
147
+ f"""No default preprocess for model {name} ({model_type}).
148
+ This can happen if the model is not finetuned on downstream datasets,
149
+ or it is not intended for direct use without finetuning.
150
+ """
151
+ )
152
+
153
+ if device == "cpu" or device == torch.device("cpu"):
154
+ model = model.float()
155
+
156
+ return model.to(device), vis_processors, txt_processors
157
+
158
+
159
+ class ModelZoo:
160
+ """
161
+ A utility class to create string representation of available model architectures and types.
162
+
163
+ >>> from unimernet.models import model_zoo
164
+ >>> # list all available models
165
+ >>> print(model_zoo)
166
+ >>> # show total number of models
167
+ >>> print(len(model_zoo))
168
+ """
169
+
170
+ def __init__(self) -> None:
171
+ self.model_zoo = {
172
+ k: list(v.PRETRAINED_MODEL_CONFIG_DICT.keys())
173
+ for k, v in registry.mapping["model_name_mapping"].items()
174
+ }
175
+
176
+ def __str__(self) -> str:
177
+ return (
178
+ "=" * 50
179
+ + "\n"
180
+ + f"{'Architectures':<30} {'Types'}\n"
181
+ + "=" * 50
182
+ + "\n"
183
+ + "\n".join(
184
+ [
185
+ f"{name:<30} {', '.join(types)}"
186
+ for name, types in self.model_zoo.items()
187
+ ]
188
+ )
189
+ )
190
+
191
+ def __iter__(self):
192
+ return iter(self.model_zoo.items())
193
+
194
+ def __len__(self):
195
+ return sum([len(v) for v in self.model_zoo.values()])
196
+
197
+
198
+ model_zoo = ModelZoo()
unimernet/models/__pycache__/__init__.cpython-310.pyc ADDED
Binary file (6.04 kB). View file
 
unimernet/models/__pycache__/base_model.cpython-310.pyc ADDED
Binary file (8.8 kB). View file
 
unimernet/models/__pycache__/clip_vit.cpython-310.pyc ADDED
Binary file (9.22 kB). View file
 
unimernet/models/__pycache__/eva_vit.cpython-310.pyc ADDED
Binary file (14.1 kB). View file
 
unimernet/models/base_model.py ADDED
@@ -0,0 +1,251 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Copyright (c) 2022, salesforce.com, inc.
3
+ All rights reserved.
4
+ SPDX-License-Identifier: BSD-3-Clause
5
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ """
7
+
8
+ import logging
9
+ import os
10
+
11
+ import numpy as np
12
+ import torch
13
+ import torch.nn as nn
14
+ from unimernet.common.dist_utils import download_cached_file, is_dist_avail_and_initialized
15
+ from unimernet.common.utils import get_abs_path, is_url
16
+ from omegaconf import OmegaConf
17
+
18
+
19
+ class BaseModel(nn.Module):
20
+ """Base class for models."""
21
+
22
+ def __init__(self):
23
+ super().__init__()
24
+
25
+ @property
26
+ def device(self):
27
+ return list(self.parameters())[0].device
28
+
29
+ def load_checkpoint(self, url_or_filename):
30
+ """
31
+ Load from a finetuned checkpoint.
32
+
33
+ This should expect no mismatch in the model keys and the checkpoint keys.
34
+ """
35
+
36
+ if is_url(url_or_filename):
37
+ cached_file = download_cached_file(
38
+ url_or_filename, check_hash=False, progress=True
39
+ )
40
+ checkpoint = torch.load(cached_file, map_location="cpu")
41
+ elif os.path.isfile(url_or_filename):
42
+ checkpoint = torch.load(url_or_filename, map_location="cpu")
43
+ else:
44
+ raise RuntimeError("checkpoint url or path is invalid")
45
+
46
+ if "model" in checkpoint.keys():
47
+ state_dict = checkpoint["model"]
48
+ else:
49
+ state_dict = checkpoint
50
+
51
+ msg = self.load_state_dict(state_dict, strict=False)
52
+
53
+ # logging.info("Missing keys {}".format(msg.missing_keys))
54
+ logging.info(f"Missing keys exist when loading '{url_or_filename}'.")
55
+ logging.info("load checkpoint from %s" % url_or_filename)
56
+
57
+ return msg
58
+
59
+ @classmethod
60
+ def from_pretrained(cls, model_type):
61
+ """
62
+ Build a pretrained model from default configuration file, specified by model_type.
63
+
64
+ Args:
65
+ - model_type (str): model type, specifying architecture and checkpoints.
66
+
67
+ Returns:
68
+ - model (nn.Module): pretrained or finetuned model, depending on the configuration.
69
+ """
70
+ model_cfg = OmegaConf.load(cls.default_config_path(model_type)).model
71
+ model = cls.from_config(model_cfg)
72
+
73
+ return model
74
+
75
+ @classmethod
76
+ def default_config_path(cls, model_type):
77
+ assert (
78
+ model_type in cls.PRETRAINED_MODEL_CONFIG_DICT
79
+ ), "Unknown model type {}".format(model_type)
80
+ return get_abs_path(cls.PRETRAINED_MODEL_CONFIG_DICT[model_type])
81
+
82
+ def load_checkpoint_from_config(self, cfg, **kwargs):
83
+ """
84
+ Load checkpoint as specified in the config file.
85
+
86
+ If load_finetuned is True, load the finetuned model; otherwise, load the pretrained model.
87
+ When loading the pretrained model, each task-specific architecture may define their
88
+ own load_from_pretrained() method.
89
+ """
90
+ load_pretrained = cfg.get("load_pretrained", True)
91
+ load_finetuned = cfg.get("load_finetuned", False)
92
+
93
+ if load_pretrained:
94
+ # load pre-trained weights
95
+ pretrain_path = cfg.get("pretrained", None)
96
+ assert pretrain_path, "Found load_finetuned is False, but pretrain_path is None."
97
+ self.load_from_pretrained(url_or_filename=pretrain_path, **kwargs)
98
+ logging.info(f"Loaded pretrained model '{pretrain_path}'.")
99
+
100
+ if load_finetuned:
101
+ finetune_path = cfg.get("finetuned", None)
102
+ assert finetune_path is not None, "Found load_finetuned is True, but finetune_path is None."
103
+ self.load_checkpoint(url_or_filename=finetune_path)
104
+ logging.info(f"Loaded finetuned model '{finetune_path}'.")
105
+
106
+ def before_evaluation(self, **kwargs):
107
+ pass
108
+
109
+ def show_n_params(self, return_str=True):
110
+ tot = 0
111
+ for p in self.parameters():
112
+ w = 1
113
+ for x in p.shape:
114
+ w *= x
115
+ tot += w
116
+ if return_str:
117
+ if tot >= 1e6:
118
+ return "{:.1f}M".format(tot / 1e6)
119
+ else:
120
+ return "{:.1f}K".format(tot / 1e3)
121
+ else:
122
+ return tot
123
+
124
+
125
+ class BaseEncoder(nn.Module):
126
+ """
127
+ Base class for primitive encoders, such as ViT, TimeSformer, etc.
128
+ """
129
+
130
+ def __init__(self):
131
+ super().__init__()
132
+
133
+ def forward_features(self, samples, **kwargs):
134
+ raise NotImplementedError
135
+
136
+ @property
137
+ def device(self):
138
+ return list(self.parameters())[0].device
139
+
140
+
141
+ class SharedQueueMixin:
142
+ @torch.no_grad()
143
+ def _dequeue_and_enqueue(self, image_feat, text_feat, idxs=None):
144
+ # gather keys before updating queue
145
+ image_feats = concat_all_gather(image_feat)
146
+ text_feats = concat_all_gather(text_feat)
147
+
148
+ batch_size = image_feats.shape[0]
149
+
150
+ ptr = int(self.queue_ptr)
151
+ assert self.queue_size % batch_size == 0 # for simplicity
152
+
153
+ # replace the keys at ptr (dequeue and enqueue)
154
+ self.image_queue[:, ptr: ptr + batch_size] = image_feats.T
155
+ self.text_queue[:, ptr: ptr + batch_size] = text_feats.T
156
+
157
+ if idxs is not None:
158
+ idxs = concat_all_gather(idxs)
159
+ self.idx_queue[:, ptr: ptr + batch_size] = idxs.T
160
+
161
+ ptr = (ptr + batch_size) % self.queue_size # move pointer
162
+ self.queue_ptr[0] = ptr
163
+
164
+
165
+ class MomentumDistilationMixin:
166
+ @torch.no_grad()
167
+ def copy_params(self):
168
+ for model_pair in self.model_pairs:
169
+ for param, param_m in zip(
170
+ model_pair[0].parameters(), model_pair[1].parameters()
171
+ ):
172
+ param_m.data.copy_(param.data) # initialize
173
+ param_m.requires_grad = False # not update by gradient
174
+
175
+ @torch.no_grad()
176
+ def _momentum_update(self):
177
+ for model_pair in self.model_pairs:
178
+ for param, param_m in zip(
179
+ model_pair[0].parameters(), model_pair[1].parameters()
180
+ ):
181
+ param_m.data = param_m.data * self.momentum + param.data * (
182
+ 1.0 - self.momentum
183
+ )
184
+
185
+
186
+ class GatherLayer(torch.autograd.Function):
187
+ """
188
+ Gather tensors from all workers with support for backward propagation:
189
+ This implementation does not cut the gradients as torch.distributed.all_gather does.
190
+ """
191
+
192
+ @staticmethod
193
+ def forward(ctx, x):
194
+ output = [
195
+ torch.zeros_like(x) for _ in range(torch.distributed.get_world_size())
196
+ ]
197
+ torch.distributed.all_gather(output, x)
198
+ return tuple(output)
199
+
200
+ @staticmethod
201
+ def backward(ctx, *grads):
202
+ all_gradients = torch.stack(grads)
203
+ torch.distributed.all_reduce(all_gradients)
204
+ return all_gradients[torch.distributed.get_rank()]
205
+
206
+
207
+ def all_gather_with_grad(tensors):
208
+ """
209
+ Performs all_gather operation on the provided tensors.
210
+ Graph remains connected for backward grad computation.
211
+ """
212
+ # Queue the gathered tensors
213
+ world_size = torch.distributed.get_world_size()
214
+ # There is no need for reduction in the single-proc case
215
+ if world_size == 1:
216
+ return tensors
217
+
218
+ # tensor_all = GatherLayer.apply(tensors)
219
+ tensor_all = GatherLayer.apply(tensors)
220
+
221
+ return torch.cat(tensor_all, dim=0)
222
+
223
+
224
+ @torch.no_grad()
225
+ def concat_all_gather(tensor):
226
+ """
227
+ Performs all_gather operation on the provided tensors.
228
+ *** Warning ***: torch.distributed.all_gather has no gradient.
229
+ """
230
+ # if use distributed training
231
+ if not is_dist_avail_and_initialized():
232
+ return tensor
233
+
234
+ tensors_gather = [
235
+ torch.ones_like(tensor) for _ in range(torch.distributed.get_world_size())
236
+ ]
237
+ torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
238
+
239
+ output = torch.cat(tensors_gather, dim=0)
240
+ return output
241
+
242
+
243
+ def tile(x, dim, n_tile):
244
+ init_dim = x.size(dim)
245
+ repeat_idx = [1] * x.dim()
246
+ repeat_idx[dim] = n_tile
247
+ x = x.repeat(*(repeat_idx))
248
+ order_index = torch.LongTensor(
249
+ np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)])
250
+ )
251
+ return torch.index_select(x, dim, order_index.to(x.device))