cfg.DATASETS.TRAIN = tuple(args.TRAIN)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = args.NUM_WORKERS
#cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = args.IMS_PER_BATCH
cfg.SOLVER.BASE_LR = args.BASE_LR # pick a good LR
cfg.SOLVER.MAX_ITER = args.MAX_ITER # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
cfg.SOLVER.STEPS = [] # do not decay learning rate
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = args.BATCH_SIZE_PER_IMAGE # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = args.NUM_CLASSES # only has one class (ballon). (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)
# NOTE: this config means the number of classes, but a few popular unofficial tutorials incorrect uses num_classes+1 here.
### BEGINNING TRAINING ###
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
### TRAINING THE MODEL ###
im = cv2.imread("./pic.jpg")
for (i=0 to args.MAX_ITER):
if (i%50==0):
#Save inference every 50 iterations
predictor = DefaultPredictor(cfg)
outputs=predictor(im)
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
imoutput=out.get_image()[:, :, ::-1]
filename=str(i)+".jpg"
cv2.imwrite(filename,imoutput)
trainer.run_step()
print("Training complete")
