This Note implements the new Detectron2 Library by Meta(or facebook). This article notebook shows training on your own custom objects for object detection.
It is worth noting that the Detectron2 library goes far beyond object detection, supporting semantic segmentation, keypoint detection, mask, and densepose. Exactly is a fantastic AI tool!
Accompanying Blog Post
We recommend that you follow along in this notebook while reading the blog post on how to train Detectron2, concurrently.
Steps Covered in this Tutorial
In this Note, we will walk through the steps required to train Detectron2 on a North American Mushroom detection dataset on roboflow, which is open source and free to use. You can also use this notebook on your own data.
To train our detector we take the following steps:
- Install Detectron2 dependencies
- Download custom Detectron2 object detection data
- Visualize Detectron2 training data
- Write our Detectron2 Training configuration
- Run Detectron2 training
- Evaluate Detectron2 performance
- Run Detectron2 inference on test images
About
Roboflow enables teams to deploy custom computer vision models quickly and accurately. Convert data from to annotation format, assess dataset health, preprocess, augment, and more. It’s free for your first 1000 source images.
Firstly,Install Detectron2 Dependencies.
all slash words represent code.
# install dependencies: (use cu101 because colab has CUDA 10.1)
!pip install -U torch==1.5 torchvision==0.6 -f https://download.pytorch.org/whl/cu101/torch_stable.html
!pip install cython pyyaml==5.1
!pip install -U ‘git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
import torch, torchvision
print(torch.__version__, torch.cuda.is_available())
!gcc — version
# opencv is pre-installed on colab
# install detectron2:
!pip install detectron2==0.1.3 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.5/index.html
output:
Collecting torch==1.5…,Collecting torchvision==0.6…Installing collected packages: torch, torchvision…
Successfully installed torch-1.5.0+cu101 torchvision-0.6.0+cu101…
Collecting pyyaml==5.1 Downloading PyYAML-5.1.tar.gz (274 kB)
Successfully installed pyyaml-5.1…(sometimes we should note our versions).
Successfully uninstalled pycocotools-2.0.4 Successfully installed pycocotools-2.0 1.5.0+cu101 True…OK we done these installations.
Successfully built fvcore iopath Installing collected packages: portalocker, yacs, iopath, mock, fvcore, detectron2 Successfully installed detectron2–0.1.3+cu101 fvcore-0.1.5.post20220512 iopath-0.1.10 mock-4.0.3 portalocker-2.5.1 yacs-0.1.8…
Next cell…
# You may need to restart your runtime prior to this, to let your installation take effect
# Some basic setup:
# Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()
# import some common libraries
import numpy as np
import cv2
import random
from google.colab.patches import cv2_imshow
# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
from detectron2.data.catalog import DatasetCatalog
Import and Register Custom Detectron2 Data
!wget https://moderncomputervision.s3.eu-west-2.amazonaws.com/North+American+Mushrooms.v1-416x416.coco.zip
!unzip -q ‘North American Mushrooms.v1–416x416.coco.zip’
!curl -L “https://public.roboflow.com/ds/6lksQTPxfb?key=FMnBpxKVpK" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip
We’ll get train,val,test .jpg and their _annotation.coco.json document like below…
from detectron2.data.datasets import register_coco_instances
register_coco_instances(“my_dataset_train”, {}, “/content/train/_annotations.coco.json”, “/content/train”)
register_coco_instances(“my_dataset_val”, {}, “/content/valid/_annotations.coco.json”, “/content/valid”)
register_coco_instances(“my_dataset_test”, {}, “/content/test/_annotations.coco.json”, “/content/test”)
#visualize training data
my_dataset_train_metadata = MetadataCatalog.get(“my_dataset_train”)
dataset_dicts = DatasetCatalog.get(“my_dataset_train”)
import random
from detectron2.utils.visualizer import Visualizer
for d in random.sample(dataset_dicts, 3):
img = cv2.imread(d[“file_name”])
visualizer = Visualizer(img[:, :, ::-1], metadata=my_dataset_train_metadata, scale=0.5)
vis = visualizer.draw_dataset_dict(d)
cv2_imshow(vis.get_image()[:, :, ::-1])
Train Custom Detectron2 Detector
#We are importing our own Trainer Module here to use the COCO validation evaluation during training. Otherwise no validation eval occurs.
from detectron2.engine import DefaultTrainer
from detectron2.evaluation import COCOEvaluator
class CocoTrainer(DefaultTrainer):
@classmethod
def build_evaluator(cls, cfg, dataset_name, output_folder=None):
if output_folder is None:
os.makedirs(“coco_eval”, exist_ok=True)
output_folder = “coco_eval”
return COCOEvaluator(dataset_name, cfg, False, output_folder)
#from .detectron2.tools.train_net import Trainer
#from detectron2.engine import DefaultTrainer
# select from modelzoo here: https://github.com/facebookresearch/detectron2/blob/master/MODEL_ZOO.md#coco-object-detection-baselines
from detectron2.config import get_cfg
#from detectron2.evaluation.coco_evaluation import COCOEvaluator
import os
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(“COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml”))
cfg.DATASETS.TRAIN = (“my_dataset_train”,)
cfg.DATASETS.TEST = (“my_dataset_val”,)
cfg.DATALOADER.NUM_WORKERS = 4
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(“COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml”) # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.001
cfg.SOLVER.WARMUP_ITERS = 1000
cfg.SOLVER.MAX_ITER = 1500 #adjust up if val mAP is still rising, adjust down if overfit
cfg.SOLVER.STEPS = (1000, 1500)
cfg.SOLVER.GAMMA = 0.05
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 64
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4 #your number of classes + 1
cfg.TEST.EVAL_PERIOD = 500
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = CocoTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
# Look at training curves in tensorboard:
%load_ext tensorboard
%tensorboard — logdir output
#test evaluation
from detectron2.data import DatasetCatalog, MetadataCatalog, build_detection_test_loader
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, “model_final.pth”)
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.85
predictor = DefaultPredictor(cfg)
evaluator = COCOEvaluator(“my_dataset_test”, cfg, False, output_dir=”./output/”)
val_loader = build_detection_test_loader(cfg, “my_dataset_test”)
inference_on_dataset(trainer.model, val_loader, evaluator)
Inference with Detectron2 Saved Weights
%ls ./output/
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, “model_final.pth”)
cfg.DATASETS.TEST = (“my_dataset_test”, )
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set the testing threshold for this model
predictor = DefaultPredictor(cfg)
test_metadata = MetadataCatalog.get(“my_dataset_test”)
from detectron2.utils.visualizer import ColorMode
import glob
for imageName in glob.glob(‘/content/test/*jpg’):
im = cv2.imread(imageName)
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
metadata=test_metadata,
scale=0.8
)
out = v.draw_instance_predictions(outputs[“instances”].to(“cpu”))
cv2_imshow(out.get_image()[:, :, ::-1])