
import torch, torchvision, os
print(torch.__version__, torch.cuda.is_available())
# assert torch.__version__.startswith("1.9")

# Some basic setup:
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import numpy as np
import os, json, random, cv2

# 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, DatasetCatalog
from detectron2.data.datasets import register_coco_instances

dir_path = os.path.dirname(os.path.realpath(__file__))

# data_root = '/home/product/dvaughn/code/detectron/data/D/'
# data_root = '../data/D/'
data_root = os.path.abspath(os.path.join(dir_path,'../../data/D/')) + '/'

img_path = data_root + 'images'
train_json_file = data_root + 'train.json'
val_json_file = data_root + 'val.json'
test_json_file = data_root + 'test.json'

register_coco_instances("damage_D_train", {}, train_json_file, img_path)
register_coco_instances("damage_D_val", {}, val_json_file, img_path)
register_coco_instances("damage_D_test", {}, test_json_file, img_path)

##########################################################
## visualize training data

'''
import random
from detectron2.utils.visualizer import Visualizer
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

my_dataset_train_metadata = MetadataCatalog.get("damage_D_train")
dataset_dicts = DatasetCatalog.get("damage_D_train")

n = 0
for d in random.sample(dataset_dicts, n):
    img = cv2.imread(d["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=my_dataset_train_metadata, scale=0.5)
    vis = visualizer.draw_dataset_dict(d)
    # M = vis.get_image()#[:, :, ::-1]
    plt.figure(figsize=(15,10))
    plt.imshow(vis.get_image())
    plt.show()
'''

##########################################################
# We are importing our own Trainer Module here to use the COCO validation evaluation during training. 
# Otherwise no validation eval occurs.

# select from modelzoo here: https://github.com/facebookresearch/detectron2/blob/master/MODEL_ZOO.md#coco-object-detection-baselines
## https://github.com/facebookresearch/detectron2/tree/main/configs
## https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md

# model_config = 'Misc/cascade_mask_rcnn_R_50_FPN_3x.yaml'
# model_config = 'COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml'
model_config = 'COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml'

from detectron2.engine import DefaultTrainer, launch, default_setup, default_argument_parser
from detectron2.evaluation import COCOEvaluator
from detectron2.config import get_cfg

class CocoTrainer(DefaultTrainer):
  @classmethod
  def build_evaluator(cls, cfg, dataset_name, output_folder=None):
    if output_folder is None:
        os.makedirs("output/coco_eval", exist_ok=True)
        output_folder = "output/coco_eval"
    return COCOEvaluator(dataset_name, cfg, False, output_folder)

##########################################################

def setup(args):
    cfg = get_cfg()
    # cfg.merge_from_file(args.config_file)
    cfg.merge_from_file(model_zoo.get_config_file(model_config))
    cfg.merge_from_list(args.opts)
    cfg.DATASETS.TRAIN = ("damage_D_train",)
    cfg.DATASETS.TEST = ("damage_D_val", "damage_D_test")
    ###
    cfg.DATALOADER.NUM_WORKERS = 4
    cfg.SOLVER.IMS_PER_BATCH = 2
    ###
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_config)
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 22
    cfg.TEST.EVAL_PERIOD = 1000
    cfg.SOLVER.CHECKPOINT_PERIOD = 2000
    cfg.SOLVER.WARMUP_ITERS = 100
    cfg.SOLVER.MAX_ITER = 200000
    cfg.SOLVER.BASE_LR = 0.00025
    cfg.SOLVER.STEPS = []
    
    cfg.INPUT.MIN_SIZE_TRAIN = (640, 672, 704, 736, 768, 800, 864, 896, 928, 960, 992, 1024)
    cfg.INPUT.MAX_SIZE_TRAIN = 1600
    cfg.INPUT.MIN_SIZE_TEST = 1024
    cfg.INPUT.MAX_SIZE_TEST = 1600

    # cfg.INPUT.MIN_SIZE_TRAIN = (800, 864, 896, 928, 960, 992, 1024, 1280, 1536, 1600)
    # cfg.INPUT.MAX_SIZE_TRAIN = 2560
    # cfg.INPUT.MIN_SIZE_TEST = 1024
    # cfg.INPUT.MAX_SIZE_TEST = 2560

    # cfg.SOLVER.REFERENCE_WORLD_SIZE = 2

    ##########################################################

    cfg.freeze()
    default_setup(cfg, args)
    return cfg

def main(args):
    cfg = setup(args)
    trainer = CocoTrainer(cfg)
    trainer.resume_or_load(resume=False)
    return trainer.train()

if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    print("Command Line Args:", args)
    launch(
        main,
        args.num_gpus,
        num_machines=args.num_machines,
        machine_rank=args.machine_rank,
        dist_url=args.dist_url,
        args=(args,),
    )

import torch, torchvision, os
print(torch.__version__, torch.cuda.is_available())
# assert torch.__version__.startswith("1.9")

# Some basic setup:
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import numpy as np
import os, json, random, cv2

# 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, DatasetCatalog
from detectron2.data.datasets import register_coco_instances

dir_path = os.path.dirname(os.path.realpath(__file__))

# data_root = '/home/product/dvaughn/code/detectron/data/D/'
# data_root = '../data/D/'
data_root = os.path.abspath(os.path.join(dir_path,'../../data/D/')) + '/'

img_path = data_root + 'images'
train_json_file = data_root + 'train.json'
val_json_file = data_root + 'val.json'
test_json_file = data_root + 'test.json'

register_coco_instances("damage_D_train", {}, train_json_file, img_path)
register_coco_instances("damage_D_val", {}, val_json_file, img_path)
register_coco_instances("damage_D_test", {}, test_json_file, img_path)

##########################################################
## visualize training data

'''
import random
from detectron2.utils.visualizer import Visualizer
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

my_dataset_train_metadata = MetadataCatalog.get("damage_D_train")
dataset_dicts = DatasetCatalog.get("damage_D_train")

n = 0
for d in random.sample(dataset_dicts, n):
    img = cv2.imread(d["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=my_dataset_train_metadata, scale=0.5)
    vis = visualizer.draw_dataset_dict(d)
    # M = vis.get_image()#[:, :, ::-1]
    plt.figure(figsize=(15,10))
    plt.imshow(vis.get_image())
    plt.show()
'''

##########################################################
# We are importing our own Trainer Module here to use the COCO validation evaluation during training. 
# Otherwise no validation eval occurs.

# select from modelzoo here: https://github.com/facebookresearch/detectron2/blob/master/MODEL_ZOO.md#coco-object-detection-baselines
## https://github.com/facebookresearch/detectron2/tree/main/configs
## https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md

# model_config = 'Misc/cascade_mask_rcnn_R_50_FPN_3x.yaml'
# model_config = 'COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml'
model_config = 'COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml'

from detectron2.engine import DefaultTrainer, launch, default_setup, default_argument_parser
from detectron2.evaluation import COCOEvaluator
from detectron2.config import get_cfg

class CocoTrainer(DefaultTrainer):
  @classmethod
  def build_evaluator(cls, cfg, dataset_name, output_folder=None):
    if output_folder is None:
        os.makedirs("output/coco_eval", exist_ok=True)
        output_folder = "output/coco_eval"
    return COCOEvaluator(dataset_name, cfg, False, output_folder)

##########################################################

def setup(args):
    cfg = get_cfg()
    # cfg.merge_from_file(args.config_file)
    cfg.merge_from_file(model_zoo.get_config_file(model_config))
    cfg.merge_from_list(args.opts)
    cfg.DATASETS.TRAIN = ("damage_D_train",)
    cfg.DATASETS.TEST = ("damage_D_val", "damage_D_test")
    ###
    cfg.DATALOADER.NUM_WORKERS = 4
    cfg.SOLVER.IMS_PER_BATCH = 2
    ###
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_config)
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 22
    cfg.TEST.EVAL_PERIOD = 1000
    cfg.SOLVER.CHECKPOINT_PERIOD = 2000
    cfg.SOLVER.WARMUP_ITERS = 100
    cfg.SOLVER.MAX_ITER = 200000
    cfg.SOLVER.BASE_LR = 0.00025
    cfg.SOLVER.STEPS = []
    
    cfg.INPUT.MIN_SIZE_TRAIN = (640, 672, 704, 736, 768, 800, 864, 896, 928, 960, 992, 1024)
    cfg.INPUT.MAX_SIZE_TRAIN = 1600
    cfg.INPUT.MIN_SIZE_TEST = 1024
    cfg.INPUT.MAX_SIZE_TEST = 1600

    # cfg.INPUT.MIN_SIZE_TRAIN = (800, 864, 896, 928, 960, 992, 1024, 1280, 1536, 1600)
    # cfg.INPUT.MAX_SIZE_TRAIN = 2560
    # cfg.INPUT.MIN_SIZE_TEST = 1024
    # cfg.INPUT.MAX_SIZE_TEST = 2560

    # cfg.SOLVER.REFERENCE_WORLD_SIZE = 2

    ##########################################################

    cfg.freeze()
    default_setup(cfg, args)
    return cfg

def main(args):
    cfg = setup(args)
    trainer = CocoTrainer(cfg)
    trainer.resume_or_load(resume=False)
    return trainer.train()

if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    print("Command Line Args:", args)
    launch(
        main,
        args.num_gpus,
        num_machines=args.num_machines,
        machine_rank=args.machine_rank,
        dist_url=args.dist_url,
        args=(args,),
    )
# assert torch.__version__.startswith("1.9") # Some basic setup: # import some common libraries # import some common detectron2 utilities # data_root = '/home/product/dvaughn/code/detectron/data/D/' # data_root = '../data/D/' ########################################################## ## visualize training data import random from detectron2.utils.visualizer import Visualizer import matplotlib matplotlib.use('TkAgg') import matplotlib.pyplot as plt  my_dataset_train_metadata = MetadataCatalog.get("damage_D_train") dataset_dicts = DatasetCatalog.get("damage_D_train")  n = 0 for d in random.sample(dataset_dicts, n):     img = cv2.imread(d["file_name"])     visualizer = Visualizer(img[:, :, ::-1], metadata=my_dataset_train_metadata, scale=0.5)     vis = visualizer.draw_dataset_dict(d)     # M = vis.get_image()#[:, :, ::-1]     plt.figure(figsize=(15,10))     plt.imshow(vis.get_image())     plt.show() ########################################################## # We are importing our own Trainer Module here to use the COCO validation evaluation during training. # Otherwise no validation eval occurs. # select from modelzoo here: https://github.com/facebookresearch/detectron2/blob/master/MODEL_ZOO.md#coco-object-detection-baselines ## https://github.com/facebookresearch/detectron2/tree/main/configs ## https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md # model_config = 'Misc/cascade_mask_rcnn_R_50_FPN_3x.yaml' # model_config = 'COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml' ########################################################## # cfg.merge_from_file(args.config_file) ### ### # cfg.INPUT.MIN_SIZE_TRAIN = (800, 864, 896, 928, 960, 992, 1024, 1280, 1536, 1600) # cfg.INPUT.MAX_SIZE_TRAIN = 2560 # cfg.INPUT.MIN_SIZE_TEST = 1024 # cfg.INPUT.MAX_SIZE_TEST = 2560 # cfg.SOLVER.REFERENCE_WORLD_SIZE = 2 ##########################################################