How to use ViTDet for customized dataset

[tairen99]

Thank you for your good work.

I am new to detectorn2 and I am trying to use detectorn2/ViTDet for the customized datasets, but I do not know where to start.

I read the README: detectorn2, but it has the training and evaluation based on the COCO and LVIS datasets.

I am wondering how I can customize it.

Thank you in advance!

📚 Documentation Issue

This issue category is for problems about existing documentation, not for asking how-to questions.

• Provide a link to an existing documentation/comment/tutorial: https://github.com/facebookresearch/detectron2/tree/main/projects/ViTDet

• How should the above documentation/comment/tutorial improve:

  1. Explanation and examples for customized datasets training
  2. Introduction to inference testing

[TranLuongBang]

same question

[josep-alana]

I posted a related question in the discussion channel: https://github.com/facebookresearch/detectron2/discussions/4461

[jshtok]

Same question. The config file 'projects/ViTDet/configs/COCO/mask_rcnn_vitdet_b_100ep.py' uses dataloader from '..common.coco_loader_lsj'. How is this related to the usual dataset registration via register_coco_instances? Do I need to prepare my own dataloader?

Thank you!

[Robotatron]

same question!

[GeorgePearse]

same question

[Kpenciler]

same question

[ken011528]

Not Sure Correct or not

For Train

Modify tools/lazyconfig_train_net.py in main(line 113)

# Custome at here
register_coco_instances("data set name", {}, "Json File Path" , "Img File Path")

cfg = LazyConfig.load(args.config_file)
cfg.dataloader.train.dataset.names = 'data set name'
cfg.dataloader.test.dataset.names = ''
cfg.model.roi_heads.num_classes = 1

cfg.dataloader.train.total_batch_size = 1
cfg.dataloader.train.num_workers = 1

cfg.train.output_dir = ''# Output Save Path
cfg.train.checkpointer.period = 1000
cfg.train.init_checkpoint = '' # ViTDet checkpoint
cfg = LazyConfig.apply_overrides(cfg, args.opts)
cfg.train.eval_period = -1 # Checkpoint eval period
print(cfg)

default_setup(cfg, args)

if args.eval_only:
    model = instantiate(cfg.model)
    model.to(cfg.train.device)
    model = create_ddp_model(model)
    DetectionCheckpointer(model).load(cfg.train.init_checkpoint)
    print(do_test(cfg, model))
else:
    do_train(args, cfg)

For Inference

import logging
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import LazyConfig, instantiate

from detectron2.data.datasets import register_coco_instances
import torch
from detectron2.engine.defaults import DefaultPredictor
from detectron2.utils.visualizer import ColorMode, Visualizer
import detectron2.data.transforms as T
import cv2
import numpy as np
from detectron2.modeling import build_model
from detectron2.config import get_cfg
from detectron2.engine.defaults import create_ddp_model
from detectron2.checkpoint import DetectionCheckpointer
import time

img = cv2.imread(img)

img = T.ResizeShortestEdge(short_edge_length=800, max_size=1333).get_transform(img).apply_image(img)

img_tensor = torch.as_tensor(img.astype("float32").transpose(2,0,1))

config_file = './projects/ViTDet/configs/COCO/cascade_mask_rcnn_swin_b_in21k_50ep.py'

cfg = LazyConfig.load(config_file)

cfg.model.roi_heads.num_classes = 1

cfg.train.init_checkpoint = 'check point'

cfg.model.roi_heads.box_predictors[2].test_score_thresh = 0.7 # score threshold

model = instantiate(cfg.model)
model.to('cuda')
model.eval()
DetectionCheckpointer(model).load(cfg.train.init_checkpoint)

with torch.no_grad():
    outputs = model([{'image': img_tensor}])[0]

visoutput = None
visualizer = Visualizer(img, instance_mode = ColorMode.IMAGE)

instances = outputs['instances'].to('cpu')
visoutput = visualizer.draw_instance_predictions(predictions = instances)

## Save file
visoutput.save('output.jpg')