waxnkw / IETrans-SGG.pytorch

This is the code of ECCV 2022 (Oral) paper "Fine-Grained Scene Graph Generation with Data Transfer".
Other
89 stars 6 forks source link

Finetune on custom datasets #6

Closed Mingyuan1997 closed 1 year ago

Mingyuan1997 commented 1 year ago

Thanks for your nice work! Is there anyway to create the custom datasets and finetune the model?

waxnkw commented 1 year ago

Good question! To create a custom dataset, please: (1) create a dataset like maskrcnn_benchmark/data/datasets/visual_genome.py/VGDataset and add it at maskrcnn_benchmark/data/datasets/init.py. Note that the _loadgraph in VGDataset is very complex, you can replace it with an easy-to-use one suitable for your own annotation.

(2) provide data files including image files (e.g. xxx.jpg), annotations (e.g. datasets/vg/50/VG-SGG-with-attri.h5), image_infos (e.g. datasets/vg/image_data.json), class indices (e.g. datasets/vg/50/VG-SGG-dicts-with-attri.json).

(3) modify corresponding dataset: (a) config the dataset initialization: modify maskrcnn_benchmark/config/paths_catalog.py. (b) add an entry in DatasetCatalog.DATASETS. (c) add corresponding initialization code in DatasetCatalog.DATASETS.get() method.

(4) add config files for training. Add a yaml file at configs like sup-50.yaml. Corresponding variables should be modified, like MODEL.ROI_BOX_HEAD.NUM_CLASSES, MODEL.ROI_RELATION_HEAD.NUM_CLASSES, DATASETS.TRAIN/VAL/TEST (to the entry you add in (3b)).

(5) If your dataset uses different object categories, you need also to retrain a new detector as described in here.