bendangnuksung/temp.py

## temp.py
from detectron2.utils.logger import setup_logger
from glob import glob
setup_logger()
import copy

from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
from detectron2 import model_zoo
from detectron2.config import get_cfg
from detectron2.config import CfgNode as CN
from detectron2.data.datasets import register_coco_instances
from detectron2.engine import DefaultTrainer
from detectron2.evaluation import COCOEvaluator
from detectron2.evaluation.evaluator import DatasetEvaluator
from detectron2.data import MetadataCatalog, DatasetCatalog
import os
import pycocotools.mask as mask_util
import numpy as np
import cv2
from detectron2.data import detection_utils as utils
import torch
import detectron2.data.transforms as T
from fastcore.all import *
from ensemble_boxes import *

from detectron2.utils.logger import setup_logger
from detectron2.data import build_detection_train_loader

INPUT_MASK = 'polygon'
# INPUT_MASK = 'bitmask'

MIN_PIXELS = [175, 75, 75]
CLASS_CONFIDENCE_THRESHOLDS = [.56, .71, .27]
IOU_TH = 0.4

# import some common detectron2 utilities
# from detectron2.engine import DefaultPredictor
# from detectron2.utils.visualizer import Visualizer
# from detectron2.data import MetadataCatalog, DatasetCatalog


# Taken from https://www.kaggle.com/theoviel/competition-metric-map-iou
def precision_at(threshold, iou):
    matches = iou > threshold
    if matches.shape == (0,):
        matches = np.array([[]])
    true_positives = np.sum(matches, axis=1) == 1  # Correct objects
    false_positives = np.sum(matches, axis=0) == 0  # Missed objects
    false_negatives = np.sum(matches, axis=1) == 0  # Extra objects
    return np.sum(true_positives), np.sum(false_positives), np.sum(false_negatives)


def score_old(pred, targ):
    pred_masks = pred['instances'].pred_masks.cpu().numpy()
    enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred_masks]
    enc_targs = list(map(lambda x:x['segmentation'], targ))
    ious = mask_util.iou(enc_preds, enc_targs, [0]*len(enc_targs))
    prec = []
    for t in np.arange(0.5, 1.0, 0.05):
        tp, fp, fn = precision_at(t, ious)
        p = tp / (tp + fp + fn)
        prec.append(p)
    return np.mean(prec)


def convert_contours_to_masks(contours, image_size=(520,704)):
    masks = []
    for c in contours:
        c = np.asarray(c)
        c = np.squeeze(c)
        c = c.reshape(-1,2)
        mask = np.zeros(image_size, np.uint8)
        mask = cv2.drawContours(mask, [c], -1, 255, -1)
        mask = mask.astype(np.bool)
        masks.append(mask)
    return np.array(masks)


def nms_predictions(classes, scores, bboxes, masks,
                    iou_th, shape=(520, 704)):
#     print(len(classes))
    he, wd = shape[0], shape[1]
    boxes_list = [[x[0] / wd, x[1] / he, x[2] / wd, x[3] / he]
                  for x in bboxes]
    scores_list = [x for x in scores]
    labels_list = [x for x in classes]

    nms_bboxes, nms_scores, nms_classes = nms(
        boxes=[boxes_list],
        scores=[scores_list],
        labels=[labels_list],
        weights=None,
        iou_thr=iou_th
    )
    nms_masks = []
    for s in nms_scores:
        nms_masks.append(masks[scores.index(s)])

    return nms_classes, nms_scores, nms_masks


def get_pred_class_mode(pred_classes, wanted_pred_class_list=[0,1,2], default_pred_class=0):
    pc = pred_classes.detach().cpu().numpy()
    unique, counts = np.unique(pc, return_counts=True)
    classes = dict(zip(unique, counts))
    classes = dict(sorted(classes.items(), key=lambda item: item[1], reverse=True))
    for cls in classes.keys():
        if cls in wanted_pred_class_list:
            return cls
    return default_pred_class


def iou_and_min_pixel_process(pred):
    # pred_class = torch.mode(pred['instances'].pred_classes)[0]
    pred_class = get_pred_class_mode(pred['instances'].pred_classes)

    take = pred['instances'].scores >= CLASS_CONFIDENCE_THRESHOLDS[pred_class]

    pred_masks = pred['instances'].pred_masks[take].cpu().numpy()
    classes = pred['instances'].pred_classes[take].cpu().numpy().tolist()
    scores = pred['instances'].scores[take].cpu().numpy().tolist()
    bboxes = pred['instances'].pred_boxes[take].tensor.cpu().numpy().tolist()
    if len(classes):
        classes, scores, pred_masks = nms_predictions(classes, scores, bboxes, pred_masks, iou_th=IOU_TH)
    final_masks = []
    # used = np.zeros(im.shape[:2], dtype=int)
    used = np.zeros([520, 704], dtype=int)
    for mask in pred_masks:
        mask = mask * (1-used)
        if mask.sum() >= MIN_PIXELS[pred_class]: # skip predictions with small area
            final_masks.append(mask)

    return np.array(final_masks).astype(np.bool)


def score(pred, targ_ann, coco_annotation_segmentation='rle'):
    #print("*"*80)
    #print(targ)
    #print("*"*80)
    # pred_masks = pred['instances'].pred_masks.cpu().numpy()
    pred_masks = iou_and_min_pixel_process(pred)

    #enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred_masks]

    targ = {'annotations': targ_ann}
    #enc_targs = list(map(lambda x:x['segmentation'], targ['annotations']))
    #enc_targs = convert_contours_to_masks(enc_targs)
    #enc_targs = [mask_util.encode(np.asarray(p, order='F')) for p in enc_targs]
    #ious = mask_util.iou(enc_preds, enc_targs, [0]*len(enc_targs))

    if coco_annotation_segmentation == 'rle':
    # if annotation in rle
        enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred_masks]
        enc_targs = list(map(lambda x:x['segmentation'], targ['annotations']))

    else:
    # if annotation in polygon
        enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred_masks]
        enc_targs = list(map(lambda x:x['segmentation'], targ['annotations']))
        enc_targs = convert_contours_to_masks(enc_targs)
        enc_targs = [mask_util.encode(np.asarray(p, order='F')) for p in enc_targs]

    ious = mask_util.iou(enc_preds, enc_targs, [0]*len(enc_targs))

    prec = []
    for t in np.arange(0.5, 1.0, 0.05):
        tp, fp, fn = precision_at(t, ious)
        p = tp / (tp + fp + fn)
        prec.append(p)
    return np.mean(prec)


class MAPIOUEvaluator(DatasetEvaluator):
    def __init__(self, dataset_name):
        dataset_dicts = DatasetCatalog.get(dataset_name)
        self.annotations_cache = {item['image_id']:item['annotations'] for item in dataset_dicts}

    def reset(self):
        self.scores = []

    def process(self, inputs, outputs):
        for inp, out in zip(inputs, outputs):
            if len(out['instances']) == 0:
                self.scores.append(0)
            else:
                targ = self.annotations_cache[inp['image_id']]
                self.scores.append(score(out, targ))

    def evaluate(self):
        return {"MaP IoU": np.mean(self.scores)}


SIZE = 704

HEIGHT = 520
WIDTH = 704

import random


# def get_random_transform_list(height, width):
def get_random_transform_list():

    transform_list = [
                T.RandomFlip(prob=0.5, horizontal=False, vertical=True),
                T.RandomFlip(prob=0.5, horizontal=True, vertical=False),
        ]

    if random.randint(0, 1):
        if random.randint(0, 3) == 0:
            transform_list.append(T.RandomRotation(angle=[1.0, 180.0]))

        if random.randint(0, 3) == 0:
            r_crop_size_height = random.randint(HEIGHT - 120, HEIGHT - 60)
            r_crop_size_width = random.randint(WIDTH - 140, WIDTH - 80)
            transform_list.append(T.RandomCrop("absolute", (r_crop_size_height, r_crop_size_width)))

        if random.randint(0, 3) == 0:
            r_n = random.randint(0, 3)
            if r_n == 0:
                transform_list.append(T.RandomBrightness(0.8, 1.2))

            if r_n == 1:
                transform_list.append(T.RandomContrast(0.9, 1.2))

            if r_n == 2:
                transform_list.append(T.RandomSaturation(0.9, 1.1))

            if r_n == 3:
                transform_list.append(T.RandomLighting(0.9))

        if random.randint(0, 3) == 0:
            transform_list.append(T.ResizeScale(min_scale=0.5, max_scale=0.85, target_height=HEIGHT, target_width=WIDTH))

        transform_list.append(T.Resize([HEIGHT, WIDTH]))

    return transform_list


def custom_mapper(dataset_dict):
    # global INPUT_MASK
    # print("*"*60)
    # print(dataset_dict.keys())

    dataset_dict = copy.deepcopy(dataset_dict)
    image = utils.read_image(dataset_dict["file_name"], format="BGR")
    transform_list = get_random_transform_list()

    image, transforms = T.apply_transform_gens(transform_list, image)
    dataset_dict["image"] = torch.as_tensor(image.transpose(2, 0, 1).astype("float32"))

    # annos = [
    #     utils.transform_instance_annotations(obj, transforms, image.shape[:2])
    #     for obj in dataset_dict.pop("annotations")
    #     if obj.get("iscrowd", 0) == 0
    # ]

    annos = []
    for obj in dataset_dict["annotations"]:
        ann = utils.transform_instance_annotations(obj, transforms, image.shape[:2])
        if obj.get("iscrowd", 0) == 0 and ann["segmentation"].max() > 0:
            annos.append(ann)

    instances = utils.annotations_to_instances(annos, image.shape[:2], mask_format='bitmask')
    dataset_dict["instances"] = utils.filter_empty_instances(instances)
    # print(dataset_dict["instances"])
    # exit()
    return dataset_dict


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)

    # @classmethod
    # def build_train_loader(cls, cfg):
    #     return build_detection_train_loader(cfg, mapper=custom_mapper)

    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
        return MAPIOUEvaluator(dataset_name)


def setup_config(args):
    global INPUT_MASK
    n_images = len(glob(os.path.join(args['train_image'], 'train/*.png')))
    steps_per_epoch = n_images // eval(args['batch'])

    cfg = get_cfg()
    try:
        cfg.merge_from_file(model_zoo.get_config_file(args['config_file']))
    except:
        cfg.merge_from_file(args['config_file'])

    PERIOD = 400
    cfg.DATASETS.TRAIN = ("car_parts_train",)
    cfg.DATASETS.TEST = ("car_parts_val",)
    cfg.DATASETS.VALIDATION = ("car_parts_val",)
    cfg.TEST.EVAL_PERIOD = PERIOD
    #cfg.TEST.DETECTIONS_PER_IMAGE = 100 # added
    cfg.DATALOADER.NUM_WORKERS = eval(args['worker'])
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(args['config_file'])
    cfg.SOLVER.IMS_PER_BATCH = eval(args['batch'])
    cfg.SOLVER.CHECKPOINT_PERIOD = PERIOD
    cfg.SOLVER.BASE_LR = eval(args['learning_rate'])
    cfg.SOLVER.MAX_ITER = (steps_per_epoch * eval(args['epoch']))
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = eval(args['size'])
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = eval(args['classes'])  # only has one class (ballon)
    cfg.MODEL.MASK_ON = True
    cfg.OUTPUT_DIR = args['output']
    cfg.INPUT.RANDOM_FLIP = "horizontal"

    # cfg.INPUT.RANDOM_FLIP = "vertical"
    cfg.INPUT.CROP = CN({"ENABLED": True})
    # cfg.INPUT.CROP.TYPE = "absolute"
    # cfg.INPUT.CROP.SIZE = [460, 620]
    cfg.INPUT.CROP.TYPE = "relative"
    cfg.INPUT.CROP.SIZE = [0.8, 0.8]

    # cfg.INPUT.RANDOM_FLIP = "vertical"
    # cfg.INPUT.CROP = CN({'ENABLED': True, 'TYPE': 'relative_range', 'SIZE': [0.8, 0.75]})

    # size change
    # cfg.INPUT.MAX_SIZE_TEST = SIZE
    # cfg.INPUT.MAX_SIZE_TRAIN = SIZE
    # cfg.INPUT.MIN_SIZE_TEST = SIZE
    # cfg.INPUT.MIN_SIZE_TRAIN = (SIZE)

    setup_logger(args['output'])
    cfg.INPUT.MASK_FORMAT = 'bitmask'
    if cfg.INPUT.MASK_FORMAT == 'bitmask':
        INPUT_MASK = 'rle'

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

    # pretrained_model_path = '/camcom/ben/models/casx152/c10/4/top/model_0036499.pth'
    # cfg.MODEL.WEIGHTS = pretrained_model_path

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

    if True:
        cfg.SOLVER.WARMUP_ITERS = 750
        # cfg.SOLVER.WEIGHT_DECAY = 0.0002
        cfg.SOLVER.MAX_ITER = 75000
        # cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256

    #########################################################################
    #########################################################################
    #########################################################################
    #########################################################################
    # finetune existing model
    # if True:
    #     cfg.SOLVER.BASE_LR = 0.000005
    #     cfg.SOLVER.WARMUP_ITERS = 0
    #     cfg.SOLVER.WEIGHT_DECAY = 0.0000005
    #     cfg.MODEL.WEIGHTS = "/camcom/ben/models/casx152/c10/4/top/model_0036249.pth"
    #     cfg.TEST.EVAL_PERIOD = 200
    #     cfg.SOLVER.CHECKPOINT_PERIOD = 200

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

    # extras
    # cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
    #cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
    # cfg.SOLVER.BASE_LR = 0.006
    return cfg


def train(args):
    register_coco_instances("car_parts_train", {}, args['train_json'], args['train_image'])
    register_coco_instances("car_parts_val", {}, args['val_json'], args['val_image'])

    cfg = setup_config(args)

    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
    trainer = CocoTrainer(cfg)

    resume = eval(args['resume'])
    trainer.resume_or_load(resume=resume)
    print("*" * 20)
    print(cfg)
    print("*" * 20)
    trainer.train()

    print("*"*50)
    print("Finished Training")
    print("*" * 50)

    evaluator = COCOEvaluator("car_parts_val", cfg, False, output_dir="./output/")
    val_loader = build_detection_test_loader(cfg, "car_parts_val")
    print(inference_on_dataset(trainer.model, val_loader, evaluator))


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description="Detectron2 training parameters")

    parser.add_argument("-tj", "--train_json", help="Path to train.json path")
    parser.add_argument("-ti", "--train_image", help="Path to train images path")

    parser.add_argument("-vj", "--val_json", help="Path to train.json path")
    parser.add_argument("-vi", "--val_image", help="Path to train images path")

    parser.add_argument("-e", "--epoch", help="Number of Epochs", default='100')
    parser.add_argument("-b", "--batch", help="Number of batch images", default='6')
    parser.add_argument("-lr", "--learning_rate", help="Number of batch images", default='0.0005')
    parser.add_argument("-s", "--size", help="batch image size", default='512')
    parser.add_argument("-c", "--classes", help="number of class labels", default='3')
    parser.add_argument("-r", "--resume", help="resume from last training", default='False')

    parser.add_argument("-w", "--worker", help="Number of workers", default='6')

    parser.add_argument("-o", "--output", help="Output Directory path", default='./slawek')

    parser.add_argument("-cfg", "--config_file", help="Model config YAML file", default='COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml')

    args = parser.parse_args()
    args = vars(args)

    train(args)
	from detectron2.utils.logger import setup_logger
	from glob import glob
	setup_logger()
	import copy

	from detectron2.evaluation import COCOEvaluator, inference_on_dataset
	from detectron2.data import build_detection_test_loader
	from detectron2 import model_zoo
	from detectron2.config import get_cfg
	from detectron2.config import CfgNode as CN
	from detectron2.data.datasets import register_coco_instances
	from detectron2.engine import DefaultTrainer
	from detectron2.evaluation import COCOEvaluator
	from detectron2.evaluation.evaluator import DatasetEvaluator
	from detectron2.data import MetadataCatalog, DatasetCatalog
	import os
	import pycocotools.mask as mask_util
	import numpy as np
	import cv2
	from detectron2.data import detection_utils as utils
	import torch
	import detectron2.data.transforms as T
	from fastcore.all import *
	from ensemble_boxes import *

	from detectron2.utils.logger import setup_logger
	from detectron2.data import build_detection_train_loader

	INPUT_MASK = 'polygon'
	# INPUT_MASK = 'bitmask'

	MIN_PIXELS = [175, 75, 75]
	CLASS_CONFIDENCE_THRESHOLDS = [.56, .71, .27]
	IOU_TH = 0.4

	# import some common detectron2 utilities
	# from detectron2.engine import DefaultPredictor
	# from detectron2.utils.visualizer import Visualizer
	# from detectron2.data import MetadataCatalog, DatasetCatalog


	# Taken from https://www.kaggle.com/theoviel/competition-metric-map-iou
	def precision_at(threshold, iou):
	matches = iou > threshold
	if matches.shape == (0,):
	matches = np.array([[]])
	true_positives = np.sum(matches, axis=1) == 1 # Correct objects
	false_positives = np.sum(matches, axis=0) == 0 # Missed objects
	false_negatives = np.sum(matches, axis=1) == 0 # Extra objects
	return np.sum(true_positives), np.sum(false_positives), np.sum(false_negatives)


	def score_old(pred, targ):
	pred_masks = pred['instances'].pred_masks.cpu().numpy()
	enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred_masks]
	enc_targs = list(map(lambda x:x['segmentation'], targ))
	ious = mask_util.iou(enc_preds, enc_targs, [0]*len(enc_targs))
	prec = []
	for t in np.arange(0.5, 1.0, 0.05):
	tp, fp, fn = precision_at(t, ious)
	p = tp / (tp + fp + fn)
	prec.append(p)
	return np.mean(prec)


	def convert_contours_to_masks(contours, image_size=(520,704)):
	masks = []
	for c in contours:
	c = np.asarray(c)
	c = np.squeeze(c)
	c = c.reshape(-1,2)
	mask = np.zeros(image_size, np.uint8)
	mask = cv2.drawContours(mask, [c], -1, 255, -1)
	mask = mask.astype(np.bool)
	masks.append(mask)
	return np.array(masks)


	def nms_predictions(classes, scores, bboxes, masks,
	iou_th, shape=(520, 704)):
	# print(len(classes))
	he, wd = shape[0], shape[1]
	boxes_list = [[x[0] / wd, x[1] / he, x[2] / wd, x[3] / he]
	for x in bboxes]
	scores_list = [x for x in scores]
	labels_list = [x for x in classes]

	nms_bboxes, nms_scores, nms_classes = nms(
	boxes=[boxes_list],
	scores=[scores_list],
	labels=[labels_list],
	weights=None,
	iou_thr=iou_th
	)
	nms_masks = []
	for s in nms_scores:
	nms_masks.append(masks[scores.index(s)])

	return nms_classes, nms_scores, nms_masks


	def get_pred_class_mode(pred_classes, wanted_pred_class_list=[0,1,2], default_pred_class=0):
	pc = pred_classes.detach().cpu().numpy()
	unique, counts = np.unique(pc, return_counts=True)
	classes = dict(zip(unique, counts))
	classes = dict(sorted(classes.items(), key=lambda item: item[1], reverse=True))
	for cls in classes.keys():
	if cls in wanted_pred_class_list:
	return cls
	return default_pred_class


	def iou_and_min_pixel_process(pred):
	# pred_class = torch.mode(pred['instances'].pred_classes)[0]
	pred_class = get_pred_class_mode(pred['instances'].pred_classes)

	take = pred['instances'].scores >= CLASS_CONFIDENCE_THRESHOLDS[pred_class]

	pred_masks = pred['instances'].pred_masks[take].cpu().numpy()
	classes = pred['instances'].pred_classes[take].cpu().numpy().tolist()
	scores = pred['instances'].scores[take].cpu().numpy().tolist()
	bboxes = pred['instances'].pred_boxes[take].tensor.cpu().numpy().tolist()
	if len(classes):
	classes, scores, pred_masks = nms_predictions(classes, scores, bboxes, pred_masks, iou_th=IOU_TH)
	final_masks = []
	# used = np.zeros(im.shape[:2], dtype=int)
	used = np.zeros([520, 704], dtype=int)
	for mask in pred_masks:
	mask = mask * (1-used)
	if mask.sum() >= MIN_PIXELS[pred_class]: # skip predictions with small area
	final_masks.append(mask)

	return np.array(final_masks).astype(np.bool)


	def score(pred, targ_ann, coco_annotation_segmentation='rle'):
	#print(""80)
	#print(targ)
	#print(""80)
	# pred_masks = pred['instances'].pred_masks.cpu().numpy()
	pred_masks = iou_and_min_pixel_process(pred)

	#enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred_masks]

	targ = {'annotations': targ_ann}
	#enc_targs = list(map(lambda x:x['segmentation'], targ['annotations']))
	#enc_targs = convert_contours_to_masks(enc_targs)
	#enc_targs = [mask_util.encode(np.asarray(p, order='F')) for p in enc_targs]
	#ious = mask_util.iou(enc_preds, enc_targs, [0]*len(enc_targs))

	if coco_annotation_segmentation == 'rle':
	# if annotation in rle
	enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred_masks]
	enc_targs = list(map(lambda x:x['segmentation'], targ['annotations']))

	else:
	# if annotation in polygon
	enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred_masks]
	enc_targs = list(map(lambda x:x['segmentation'], targ['annotations']))
	enc_targs = convert_contours_to_masks(enc_targs)
	enc_targs = [mask_util.encode(np.asarray(p, order='F')) for p in enc_targs]

	ious = mask_util.iou(enc_preds, enc_targs, [0]*len(enc_targs))

	prec = []
	for t in np.arange(0.5, 1.0, 0.05):
	tp, fp, fn = precision_at(t, ious)
	p = tp / (tp + fp + fn)
	prec.append(p)
	return np.mean(prec)


	class MAPIOUEvaluator(DatasetEvaluator):
	def __init__(self, dataset_name):
	dataset_dicts = DatasetCatalog.get(dataset_name)
	self.annotations_cache = {item['image_id']:item['annotations'] for item in dataset_dicts}

	def reset(self):
	self.scores = []

	def process(self, inputs, outputs):
	for inp, out in zip(inputs, outputs):
	if len(out['instances']) == 0:
	self.scores.append(0)
	else:
	targ = self.annotations_cache[inp['image_id']]
	self.scores.append(score(out, targ))

	def evaluate(self):
	return {"MaP IoU": np.mean(self.scores)}


	SIZE = 704

	HEIGHT = 520
	WIDTH = 704

	import random


	# def get_random_transform_list(height, width):
	def get_random_transform_list():

	transform_list = [
	T.RandomFlip(prob=0.5, horizontal=False, vertical=True),
	T.RandomFlip(prob=0.5, horizontal=True, vertical=False),
	]

	if random.randint(0, 1):
	if random.randint(0, 3) == 0:
	transform_list.append(T.RandomRotation(angle=[1.0, 180.0]))

	if random.randint(0, 3) == 0:
	r_crop_size_height = random.randint(HEIGHT - 120, HEIGHT - 60)
	r_crop_size_width = random.randint(WIDTH - 140, WIDTH - 80)
	transform_list.append(T.RandomCrop("absolute", (r_crop_size_height, r_crop_size_width)))

	if random.randint(0, 3) == 0:
	r_n = random.randint(0, 3)
	if r_n == 0:
	transform_list.append(T.RandomBrightness(0.8, 1.2))

	if r_n == 1:
	transform_list.append(T.RandomContrast(0.9, 1.2))

	if r_n == 2:
	transform_list.append(T.RandomSaturation(0.9, 1.1))

	if r_n == 3:
	transform_list.append(T.RandomLighting(0.9))

	if random.randint(0, 3) == 0:
	transform_list.append(T.ResizeScale(min_scale=0.5, max_scale=0.85, target_height=HEIGHT, target_width=WIDTH))

	transform_list.append(T.Resize([HEIGHT, WIDTH]))

	return transform_list


	def custom_mapper(dataset_dict):
	# global INPUT_MASK
	# print(""60)
	# print(dataset_dict.keys())

	dataset_dict = copy.deepcopy(dataset_dict)
	image = utils.read_image(dataset_dict["file_name"], format="BGR")
	transform_list = get_random_transform_list()

	image, transforms = T.apply_transform_gens(transform_list, image)
	dataset_dict["image"] = torch.as_tensor(image.transpose(2, 0, 1).astype("float32"))

	# annos = [
	# utils.transform_instance_annotations(obj, transforms, image.shape[:2])
	# for obj in dataset_dict.pop("annotations")
	# if obj.get("iscrowd", 0) == 0
	# ]

	annos = []
	for obj in dataset_dict["annotations"]:
	ann = utils.transform_instance_annotations(obj, transforms, image.shape[:2])
	if obj.get("iscrowd", 0) == 0 and ann["segmentation"].max() > 0:
	annos.append(ann)

	instances = utils.annotations_to_instances(annos, image.shape[:2], mask_format='bitmask')
	dataset_dict["instances"] = utils.filter_empty_instances(instances)
	# print(dataset_dict["instances"])
	# exit()
	return dataset_dict


	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)

	# @classmethod
	# def build_train_loader(cls, cfg):
	# return build_detection_train_loader(cfg, mapper=custom_mapper)

	@classmethod
	def build_evaluator(cls, cfg, dataset_name, output_folder=None):
	return MAPIOUEvaluator(dataset_name)


	def setup_config(args):
	global INPUT_MASK
	n_images = len(glob(os.path.join(args['train_image'], 'train/*.png')))
	steps_per_epoch = n_images // eval(args['batch'])

	cfg = get_cfg()
	try:
	cfg.merge_from_file(model_zoo.get_config_file(args['config_file']))
	except:
	cfg.merge_from_file(args['config_file'])

	PERIOD = 400
	cfg.DATASETS.TRAIN = ("car_parts_train",)
	cfg.DATASETS.TEST = ("car_parts_val",)
	cfg.DATASETS.VALIDATION = ("car_parts_val",)
	cfg.TEST.EVAL_PERIOD = PERIOD
	#cfg.TEST.DETECTIONS_PER_IMAGE = 100 # added
	cfg.DATALOADER.NUM_WORKERS = eval(args['worker'])
	cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(args['config_file'])
	cfg.SOLVER.IMS_PER_BATCH = eval(args['batch'])
	cfg.SOLVER.CHECKPOINT_PERIOD = PERIOD
	cfg.SOLVER.BASE_LR = eval(args['learning_rate'])
	cfg.SOLVER.MAX_ITER = (steps_per_epoch * eval(args['epoch']))
	cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = eval(args['size'])
	cfg.MODEL.ROI_HEADS.NUM_CLASSES = eval(args['classes']) # only has one class (ballon)
	cfg.MODEL.MASK_ON = True
	cfg.OUTPUT_DIR = args['output']
	cfg.INPUT.RANDOM_FLIP = "horizontal"

	# cfg.INPUT.RANDOM_FLIP = "vertical"
	cfg.INPUT.CROP = CN({"ENABLED": True})
	# cfg.INPUT.CROP.TYPE = "absolute"
	# cfg.INPUT.CROP.SIZE = [460, 620]
	cfg.INPUT.CROP.TYPE = "relative"
	cfg.INPUT.CROP.SIZE = [0.8, 0.8]

	# cfg.INPUT.RANDOM_FLIP = "vertical"
	# cfg.INPUT.CROP = CN({'ENABLED': True, 'TYPE': 'relative_range', 'SIZE': [0.8, 0.75]})

	# size change
	# cfg.INPUT.MAX_SIZE_TEST = SIZE
	# cfg.INPUT.MAX_SIZE_TRAIN = SIZE
	# cfg.INPUT.MIN_SIZE_TEST = SIZE
	# cfg.INPUT.MIN_SIZE_TRAIN = (SIZE)

	setup_logger(args['output'])
	cfg.INPUT.MASK_FORMAT = 'bitmask'
	if cfg.INPUT.MASK_FORMAT == 'bitmask':
	INPUT_MASK = 'rle'

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

	# pretrained_model_path = '/camcom/ben/models/casx152/c10/4/top/model_0036499.pth'
	# cfg.MODEL.WEIGHTS = pretrained_model_path

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

	if True:
	cfg.SOLVER.WARMUP_ITERS = 750
	# cfg.SOLVER.WEIGHT_DECAY = 0.0002
	cfg.SOLVER.MAX_ITER = 75000
	# cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256

	#########################################################################
	#########################################################################
	#########################################################################
	#########################################################################
	# finetune existing model
	# if True:
	# cfg.SOLVER.BASE_LR = 0.000005
	# cfg.SOLVER.WARMUP_ITERS = 0
	# cfg.SOLVER.WEIGHT_DECAY = 0.0000005
	# cfg.MODEL.WEIGHTS = "/camcom/ben/models/casx152/c10/4/top/model_0036249.pth"
	# cfg.TEST.EVAL_PERIOD = 200
	# cfg.SOLVER.CHECKPOINT_PERIOD = 200

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

	# extras
	# cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
	#cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
	# cfg.SOLVER.BASE_LR = 0.006
	return cfg


	def train(args):
	register_coco_instances("car_parts_train", {}, args['train_json'], args['train_image'])
	register_coco_instances("car_parts_val", {}, args['val_json'], args['val_image'])

	cfg = setup_config(args)

	os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
	trainer = CocoTrainer(cfg)

	resume = eval(args['resume'])
	trainer.resume_or_load(resume=resume)
	print("" 20)
	print(cfg)
	print("" 20)
	trainer.train()

	print(""50)
	print("Finished Training")
	print("" 50)

	evaluator = COCOEvaluator("car_parts_val", cfg, False, output_dir="./output/")
	val_loader = build_detection_test_loader(cfg, "car_parts_val")
	print(inference_on_dataset(trainer.model, val_loader, evaluator))


	if __name__ == '__main__':
	import argparse
	parser = argparse.ArgumentParser(description="Detectron2 training parameters")

	parser.add_argument("-tj", "--train_json", help="Path to train.json path")
	parser.add_argument("-ti", "--train_image", help="Path to train images path")

	parser.add_argument("-vj", "--val_json", help="Path to train.json path")
	parser.add_argument("-vi", "--val_image", help="Path to train images path")

	parser.add_argument("-e", "--epoch", help="Number of Epochs", default='100')
	parser.add_argument("-b", "--batch", help="Number of batch images", default='6')
	parser.add_argument("-lr", "--learning_rate", help="Number of batch images", default='0.0005')
	parser.add_argument("-s", "--size", help="batch image size", default='512')
	parser.add_argument("-c", "--classes", help="number of class labels", default='3')
	parser.add_argument("-r", "--resume", help="resume from last training", default='False')

	parser.add_argument("-w", "--worker", help="Number of workers", default='6')

	parser.add_argument("-o", "--output", help="Output Directory path", default='./slawek')

	parser.add_argument("-cfg", "--config_file", help="Model config YAML file", default='COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml')

	args = parser.parse_args()
	args = vars(args)

	train(args)