def showGt(imgid,catNm):
dataType='val2017'
annFile = "./annotations.json"
coco=COCO(annFile)
catIds = coco.getCatIds(catNms=catNm);
imgIds = coco.getImgIds(catIds=catIds );
imgIds = coco.getImgIds(imgIds = [imgid])
img = coco.loadImgs(imgIds[np.random.randint(0,len(imgIds))])[0]
annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
anns = coco.loadAnns(annIds)
DatasetCatalog.clear()
MetadataCatalog.clear()
keypoint_names_col = [ "c1", "c2", "c3", "c4", "c5", "c6", "c7", "c8", "c9", "c10" ]
keypoint_names_row= [ 'r'+str(i) for i in range(10)]
keypoint_names=keypoint_names_col+keypoint_names_row
print(keypoint_names)
register_coco_instances("my_dataset_train", {}, "./annotations.json", "/content/VOC2007/JPEGImages")
my_md= MetadataCatalog.get("my_dataset_train").set(
thing_classes=["table","r"],
#keypoint_names=keypoint_names,
#keypoint_connection_rules=[]
#,keypoint_flip_map=[]
)
f = '/content/VOC2007/JPEGImages/'+img['file_name']
img = cv2.imread(f)
visualizer = Visualizer(img[:, :, ::-1], metadata=my_md)
for an in anns:
an["bbox_mode"]=BoxMode.XYWH_ABS
d = {'annotations':anns}
out = visualizer.draw_dataset_dict(d)
return out.get_image()[:, :, ::-1]
#cv2_imshow()
def initData():
train_keys = ["icdar", "bank"]
DatasetCatalog.clear()
for d in ["trainval"]:
DatasetCatalog.register("table_" + d,
lambda d=d: get_icdar_dicts(train_keys))
MetadataCatalog.get("table_" + d).set(thing_classes=["border"])
def register_datasets_type(path_to_dataset_dir,
use_direction_classes,
validation_size=0.25,
test_size=0.25,
seed=42,
clear=True):
if clear:
DatasetCatalog.clear()
dataset_names = ["train", "test"
] if validation_size == 0.0 else ["train", "val", "test"]
configs = get_dataset_configs(path_to_dataset_dir, validation_size,
test_size, seed)
dataset_configs = dict(zip(dataset_names, configs))
if use_direction_classes:
thing_classes = list(DIRECTION_CLASSES)
else:
thing_classes = ["zebra_fish"]
for name in dataset_names:
DatasetCatalog.register(
get_name_with_prefix(name, use_direction_classes),
lambda: config_to_dataset(path_to_dataset_dir, dataset_configs[
name], use_direction_classes))
MetadataCatalog.get(get_name_with_prefix(
name, use_direction_classes)).set(thing_classes=thing_classes)
def register_data(path, prefix='yeast_cells_'):
"""Register all data sets as {prefix}_{setname}, i.e. yeast_cells_train"""
assert (
os.path.exists(f'{path}/labels.umsgpack')
or os.path.exists(f'{path}/labels.json')), (
"Labels not found, ensure either labels.umsgpack or labels.json "
f"exists at {path}.")
if os.path.exists(f'{path}/labels.umsgpack'):
with open(f'{path}/labels.umsgpack', 'rb') as f:
labels = umsgpack.unpack(f, encoding="utf-8")
else:
with open(f'{path}/labels.json', 'r') as f:
labels = json.load(f)
labels = validate_labels(labels, path)
DatasetCatalog.clear()
for label in labels:
DatasetCatalog.register(f"{prefix}{label}",
lambda label_=label: labels[label_])
MetadataCatalog.get(f"{prefix}{label}").set(
thing_classes=["yeast_cell"])
# yeast_cells_metadata = MetadataCatalog.get(f"{prefix}train")
return labels
def blendproc_bbox_trainer(self):
##### ap calculation
DatasetCatalog.clear()
MetadataCatalog.clear()
# DatasetCatalog.remove(val_coco_ann_path)
trainer = Detectron2BaseEvalTrainer(
train_coco_ann_path=self.coco_ann_path,
train_img_dir=self.img_dir,
val_coco_ann_path=self.val_coco_ann_path,
val_img_dir=self.val_img_path,
images_per_batch=1,
base_lr=0.003,
max_iter=1000,
output_dir=self.module_save_path,
min_size_train=self.min_size_train,
max_size_train=self.max_size_train,
checkpoint_period=200,
eval_period=100,
save_vis=True,
latest_vis_only=True,
# trainer_constructor=COCO_BBox_EvalTrainer_BlenderProc, #custom creation with mapper = none
trainer_constructor=
COCO_BBox_EvalTrainer, #custom creation with mapper = deafult mapper
model_type='COCO-Detection',
model_name=self.model_name,
)
trainer.train(resume=False)
def register_datasets(dataset_storage: dict):
DatasetCatalog.clear()
for name in dataset_storage.keys():
for mode in MODES:
ds = dataset_storage[name][mode]
ds_name = f'clash_{name}_{mode}'
DatasetCatalog.register(ds_name, lambda x=None: ds['dicts'])
MetadataCatalog.get(ds_name).set(thing_classes=ds['unit_classes'])
def get_predictor(config, model_weights, dataset, properties, things):
cfg = get_cfg()
add_property_config(cfg)
cfg.merge_from_file(os.path.join(dirname, config))
cfg.MODEL.WEIGHTS = model_weights
DatasetCatalog.clear()
DatasetCatalog.register(dataset, get_dicts)
MetadataCatalog.get(dataset).set(thing_classes=things)
MetadataCatalog.get(dataset).set(property_classes=properties)
cfg.DATASETS.TEST = (dataset,)
return DefaultPredictor(cfg)
def register_train_test(dataset_dict,metadata,train_ratio=0.66,seed=0):
DatasetCatalog.clear()
MetadataCatalog.clear()
nb_input= len(dataset_dict)
x=np.arange(nb_input)
random.Random(seed).shuffle(x)
idx_split = int(len(x) * train_ratio)
DatasetCatalog.register("datasetTrain", my_dataset_function({"images":np.array(dataset_dict)[x[:idx_split]],"metadata":metadata}))
DatasetCatalog.register("datasetTest", my_dataset_function({"images":np.array(dataset_dict)[x[idx_split:]],"metadata":metadata}))
MetadataCatalog.get("datasetTrain").stuff_classes = [v for k,v in metadata["category_names"].items()]
MetadataCatalog.get("datasetTest").stuff_classes = [v for k,v in metadata["category_names"].items()]
"""if ignoreValue is not None:
def train(config_file, image_path, annot_file, out_filename="results.txt"):
train_split = lambda: get_AICity_dataset(
image_path, annot_file, mode='first', is_train=True)
test_split = lambda: get_AICity_dataset(
image_path, annot_file, mode='first')
DatasetCatalog.clear()
DatasetCatalog.register("ai_city_train", train_split)
MetadataCatalog.get('ai_city_train').set(
thing_classes=[k for k in thing_classes.keys()])
DatasetCatalog.register("ai_city_test", test_split)
MetadataCatalog.get('ai_city_test').set(
thing_classes=[k for k in thing_classes.keys()])
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(config_file))
cfg.DATASETS.TRAIN = ("ai_city_train", )
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 4
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 = 8
cfg.SOLVER.BASE_LR = 0.0001
cfg.SOLVER.MAX_ITER = 100
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.MODEL.RETINANET.NUM_CLASSES = 1
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
evaluator = COCOEvaluator("ai_city_test",
cfg,
False,
output_dir="./output/")
val_loader = build_detection_test_loader(cfg, "ai_city_test")
inference_on_dataset(trainer.model, val_loader, evaluator)
ims_from_test = [annot['file_name'] for annot in test_split()]
det_bb = inference(config_file,
ims_from_test,
weight="./output/model_final.pth",
save_results=True,
out_filename=out_filename)
return det_bb
def evaluate():
train_keys = ["bank"]
DatasetCatalog.clear()
MetadataCatalog.clear()
keypoint_names = [
"c1", "c2", "c3", "c4", "c5", "c6", "c7", "c8", "c9", "c10"
]
keypoint_connection_rules = [[1, 2, [255, 175, 100]],
[2, 3, [255, 175, 100]],
[3, 4, [255, 175, 100]],
[4, 5, [255, 175, 100]],
[5, 6, [255, 175, 100]],
[6, 7, [255, 175, 100]],
[7, 8, [255, 175, 100]],
[8, 9, [255, 175, 100]],
[9, 10, [255, 175, 100]]]
keypoint_flip_map = [('c1', 'c10'), ('c2', 'c9'), ('c3', 'c8'),
('c4', 'c7'), ('c5', 'c6')]
for d in ["testval"]:
#DatasetCatalog.register("table_testval", lambda d=d: get_icdar_dicts(train_keys,'testjson'))
register_coco_instances("table_testval", {}, "./annotations.json",
"/content/VOC2007/JPEGImages")
MetadataCatalog.get("table_testval").set(
thing_classes=["table", "r"],
keypoint_names=keypoint_names,
keypoint_connection_rules=keypoint_connection_rules,
keypoint_flip_map=keypoint_flip_map)
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
cfg.DATALOADER.NUM_WORKERS = 2
cfg.DATASETS.TRAIN = ("table_testval", )
cfg.DATASETS.TEST = ("table_testval", )
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.MODEL.DEVICE = "cpu"
cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR
cfg.SOLVER.MAX_ITER = 300 # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2 # only has one class (ballon)
cfg.MODEL.WEIGHTS = '/content/keypoints/workdir/savedmodel/model_final.pth'
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3 # set the testing threshold for this model
cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS = 10
model = build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=True)
do_test(cfg, model)
def train(output, iou=None, nms=None, rpn=None):
batch_size = 2
DatasetCatalog.clear()
register_kitti_mots_dataset("datasets/KITTI-MOTS/training/image_02",
"datasets/KITTI-MOTS/instances_txt",
("kitti_mots_train", "kitti_mots_test"),
image_extension="png")
cfg_file = "Cityscapes/mask_rcnn_R_50_FPN.yaml"
output_dir = output
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(cfg_file))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(cfg_file)
cfg.SEED = 42
cfg.DATASETS.TRAIN = ("kitti_mots_train",)
cfg.DATASETS.TEST = ("kitti_mots_test",)
cfg.DATALOADER.NUM_WORKERS = 4
cfg.SOLVER.IMS_PER_BATCH = batch_size
cfg.SOLVER.BASE_LR = 0.0002 * batch_size / 16 # pick a good LR
cfg.SOLVER.MAX_ITER = 7500
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.OUTPUT_DIR = output_dir
if iou is not None:
cfg.MODEL.RPN.IOU_THRESHOLDS = [iou[0], iou[1]]
if nms is not None:
cfg.MODEL.RPN.NMS_THRESH = nms
if rpn is not None:
cfg.MODEL.RPN.PRE_NMS_TOPK_TRAIN = rpn[0]
cfg.MODEL.RPN.PRE_NMS_TOPK_TEST = rpn[1]
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
val_loss = ValidationLoss(cfg)
trainer.register_hooks([val_loss])
trainer._hooks = trainer._hooks[:-2] + trainer._hooks[-2:][::-1]
trainer.resume_or_load(resume=True)
trainer.train()
evaluator = COCOEvaluator("kitti_mots_test", cfg, False, output_dir=output_dir)
trainer.test(cfg, trainer.model, evaluators=[evaluator])
plot_losses(cfg)
def register_dataset(combined_train_val=False):
"""
Registers the datasets to DatasetCatalog of Detectron2
This is to ensure that Detectron can load our dataset from it's configs
:param combined_train_val: Optional parameter for combining train and validation
set
:return: None
"""
DatasetCatalog.clear()
for d in ["train_val_combined", "val", "test"
] if combined_train_val else ["train", "val", "test"]:
DatasetCatalog.register(
"indiscapes_" + d,
lambda d=d: get_indiscapes_dicts(images_root,
os.path.join(via_json_root, d)),
)
MetadataCatalog.get("indiscapes_" +
d).set(thing_classes=categories_list)
MetadataCatalog.get("indiscapes_" + d).set(evaluator_type="indiscapes")
def register_datasets(path_to_dataset_dir,
validation_size=0.25,
test_size=0.25,
seed=42,
clear=True):
if clear:
DatasetCatalog.clear()
register_datasets_type(path_to_dataset_dir,
True,
validation_size=validation_size,
test_size=test_size,
seed=seed,
clear=False)
register_datasets_type(path_to_dataset_dir,
False,
validation_size=validation_size,
test_size=test_size,
seed=seed,
clear=False)
def setup_cfg(args):
DatasetCatalog.clear()
MetadataCatalog.clear()
keypoint_names = [
"c1", "c2", "c3", "c4", "c5", "c6", "c7", "c8", "c9", "c10"
]
keypoint_connection_rules = [[1, 2, [255, 175, 100]],
[2, 3, [255, 175, 100]],
[3, 4, [255, 175, 100]],
[4, 5, [255, 175, 100]],
[5, 6, [255, 175, 100]],
[6, 7, [255, 175, 100]],
[7, 8, [255, 175, 100]],
[8, 9, [255, 175, 100]],
[9, 10, [255, 175, 100]]]
keypoint_flip_map = [('c1', 'c10'), ('c2', 'c9'), ('c3', 'c8'),
('c4', 'c7'), ('c5', 'c6')]
register_coco_instances("my_dataset_train", {},
"/content/keypointsl/annotations.json",
"/content/VOC2007/JPEGImages")
my_md = MetadataCatalog.get("my_dataset_train").set(
thing_classes=["table", "r"],
keypoint_names=keypoint_names,
keypoint_connection_rules=keypoint_connection_rules,
keypoint_flip_map=keypoint_flip_map)
cfg = get_cfg()
# cuda context is initialized before creating dataloader, so we don't fork anymore
cfg.DATALOADER.NUM_WORKERS = 0
cfg = add_export_config(cfg)
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.DATASETS.my_dataset_train = ("my_dataset_train", )
cfg.DATASETS.TEST = ("my_dataset_train", )
cfg.DATASETS.TRAIN = ("my_dataset_train", )
cfg.freeze()
if cfg.MODEL.DEVICE != "cpu":
assert TORCH_VERSION >= (
1, 5), "PyTorch>=1.5 required for GPU conversion!"
return cfg
def register_polyp_datasets(dataset_list=None):
DatasetCatalog.clear()
if dataset_list is None:
dataset_list = [
os.path.join("datasets", x) for x in os.listdir("datasets")
if os.path.islink(os.path.join("datasets", x))
]
for dataset in dataset_list:
dataset_name = str(dataset.split("/")[-1])
json_file = dataset_annots[dataset_name]
if json_file is None:
continue
else:
json_file = os.path.join(dataset, "annotations", json_file)
image_root = os.path.join(dataset, "images")
thing_ids = [
v["id"] for k, v in polyp_categories.items() if k in ["AD", "NAD"]
]
# Mapping from the incontiguous COCO category id to an id in [0, 79]
thing_dataset_id_to_contiguous_id = {
k: i
for i, k in enumerate(thing_ids)
}
thing_classes = [
v["name"] for k, v in polyp_categories.items()
if k in ["AD", "NAD"]
]
metadata = {
"thing_classes": thing_classes,
"thing_dataset_id_to_contiguous_id":
thing_dataset_id_to_contiguous_id,
"evaluator_type": 'giana' if "test" in dataset_name else "coco",
}
register_coco_instances(dataset_name, metadata, json_file, image_root)
print("Dataset {} registered".format(dataset_name))
def predict_surface(img):
for d in ["train", "val"]:
DatasetCatalog.clear()
DatasetCatalog.register("surface_" + d,
lambda d=d: get_surface_dicts("surface/" + d))
MetadataCatalog.get("surface_" + d).set(thing_classes=["surface"])
surface_metadata = MetadataCatalog.get("surface_train")
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("surface_train", )
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 1
cfg.SOLVER.BASE_LR = 0.0025
cfg.SOLVER.MAX_ITER = 1000
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.OUTPUT_DIR = "./drive/MyDrive/surface"
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3 # set a custom testing threshold
predictor = DefaultPredictor(cfg)
dataset_dicts = "surface_img/val"
#files=sorted(glob.glob("balloon/val/*.jpg"))
teller = 1
data = []
middle = [img[0].shape[0] / 2, img[0].shape[1] / 2]
background = backgroundsubtraction(img)
cv2.imwrite("background.jpg", background)
for teller, k in enumerate(img): #go through all images
minimum = None
predictor = DefaultPredictor(cfg)
outputs = predictor(k) #predict surface of image k
"""
v = Visualizer(k[:, :, ::-1], metadata=surface_metadata, scale=0.5) #visualization of predicted surface
out=v.draw_instance_predictions(outputs["instances"].to("cpu")) #...
v=out.get_image()[:, :, ::-1] #...
cv2.imwrite("./drive/MyDrive/wkvideo/surface/"+str(teller)+".jpg",v) #...
"""
#maskoutput=outputs['instances'].pred_masks.to("cpu")[0][:2] #next paragraph: check for multiple detections and make a
#decision if the predicted surface is part of road where the cyclist ride on
maskoutput = 0
indexen = []
y = []
x = []
coordinaten = []
prev_x_min = 0
prev_x_max = k.shape[1]
if len(
outputs['instances'].pred_boxes
) == 0: #if only one road detection -> take detection as "finish road"
maskoutput = np.zeros((k.shape[0], k.shape[1]),
np.uint8) #make new mask with certain shape
else:
for index, k in enumerate(
outputs['instances'].pred_boxes.to("cpu")
): #if multiple road detections-> do some checks about position of detection(knowledge that "finish road" is mostly in middle of frame)
#go through all predictions and take boundingboxescoordinates as variable k
coordinates = k.numpy(
) #transform k from tensor object to numpy
middle = coordinates[2] - coordinates[
0] #calculate the center position of the bounding box in x direction
if middle >= prev_x_min and middle <= prev_x_max: #check if center position is between two boundaries( prev frame boundaries of bounding box position)
y.append(
coordinates[3] -
coordinates[1]) #add center of boundingbox to list
x.append(coordinates[2] - coordinates[0])
indexen.append(
index) #add index of detected boundingbox to list
coordinaten.append(
coordinates) #add coordinates of boundingbox to list
best_ind = 0 #in next paragraph: if multiple surfaces where detected: do some logic,
#to find the best matching detected surface, using the position of the last detected surface of previous image
if len(indexen) > 1:
best = None
lastone = False
for d, k in enumerate(indexen[:len(indexen) - 1]):
if abs(x[d] - x[d + 1]) > ((prev_x_max - prev_x_min) / 2):
if d == len(indexen) - 1:
lastone = True
dist = abs(x[d] - (prev_x_max - prev_x_min) / 2)
if best == None or dist < best:
best = dist
best_ind = d
if lastone:
d = len(indexen) - 1
dist = abs(x[d] - (prev_x_max - prev_x_min) / 2)
if best == None or dist < best:
best = dist
best_ind = d
indexen = [best_ind]
prev_x_min = coordinaten[best_ind][0]
prev_x_max = coordinaten[best_ind][2]
for index, k in enumerate(
outputs['instances'].pred_masks.to("cpu").numpy()):
if indexen.count(index) == 1:
maskoutput += k
maskoutput = maskoutput * 255
kernel = np.ones((9, 1), np.uint8)
maskoutput = maskoutput.astype(np.uint8)
maskoutput = cv2.dilate(maskoutput, kernel, iterations=4)
maskoutput += background
maskoutput = maskoutput.astype(np.uint8)
mask = np.ones((k.shape[0], k.shape[1]), dtype=np.uint8)
img_res = cv2.bitwise_and(mask, mask, mask=maskoutput)
data.append(img_res)
del (indexen)
del (y)
del (x)
DatasetCatalog.clear()
return data
"iscrowd": 0
}
objs.append(obj)
if objs:
record["annotations"] = objs
if objs or box_cnt == 0:
dataset_dicts.append(record)
print(
f'dataset_dicts={len(dataset_dicts)} with fold:{fold}, box_cnt:{box_cnt}'
)
return dataset_dicts
from detectron2.data import DatasetCatalog, MetadataCatalog
DatasetCatalog.clear()
for d in ["train", "val_0", "val_1", "val_2", "val_3", "val_4"]:
name = "lung_" + d
if name not in DatasetCatalog.list():
print(d)
DatasetCatalog.register(name, lambda d=d: get_anotation_dicts(d))
MetadataCatalog.get(name).set(thing_classes=["bz"])
metadata = MetadataCatalog.get("lung_train")
def save_stack_feature(train: pd.DataFrame, test: pd.DataFrame, file_path):
train.to_hdf(file_path, 'train', mode='a')
test.to_hdf(file_path, 'test', mode='a')
logger.info(f'OOF file save to :{file_path}')
return train, test
def main(cfg: DictConfig):
cur_dir = hydra.utils.get_original_cwd()
os.chdir(cur_dir)
seed_everything(cfg.data.seed)
# wandb
wandb.init(project='VinBigData-Detection')
wandb.config.update(dict(cfg.data))
wandb.config.update(dict(cfg.train))
wandb.config.update(dict(cfg.aug_kwargs_detection))
wandb.config.update(dict(cfg.classification_kwargs))
# omegaconf -> dict
rep_aug_kwargs = OmegaConf.to_container(cfg.aug_kwargs_detection)
class_name_dict = {
0: 'Aortic enlargement',
1: 'Atelectasis',
2: 'Calcification',
3: 'Cardiomegaly',
4: 'Consolidation',
5: 'ILD',
6: 'Infiltration',
7: 'Lung Opacity',
8: 'Nodule/Mass',
9: 'Other lesion',
10: 'Pleural effusion',
11: 'Pleural thickening',
12: 'Pneumothorax',
13: 'Pulmonary fibrosis',
}
# Setting --------------------------------------------------
data_dir = cfg.data.data_dir
output_dir = cfg.data.output_dir
img_size = cfg.data.img_size
backbone = cfg.data.backbone
use_class14 = cfg.data.use_class14
if os.path.exists(output_dir):
shutil.rmtree(output_dir)
if use_class14:
class_name_dict.update({14: 'No finding'})
# Register Dataset --------------------------------------------------
anno_df = pd.read_csv(os.path.join(data_dir, 'train_wbf_th0.7.csv'))
if cfg.data.use_class14:
pass
else:
anno_df = anno_df[anno_df['class_id'] != 14].reset_index(drop=True)
# Extract rad id
if cfg.data.rad_id != 'all':
anno_df = anno_df[anno_df['rad_id'].isin(cfg.data.rad_id)].reset_index()
if debug:
anno_df = anno_df.head(100)
# Split train, valid data - random
if 'valid' in cfg.data.split_method:
split_rate = float(cfg.data.split_method.split('_')[1]) / 100
unique_image_ids = anno_df['image_id'].values
unique_image_ids = np.random.RandomState(cfg.data.seed).permutation(unique_image_ids)
train_image_ids = unique_image_ids[:int(len(unique_image_ids) * (1 - split_rate))]
valid_image_ids = unique_image_ids[int(len(unique_image_ids) * (1 - split_rate)):]
DatasetCatalog.register("xray_valid", lambda d='valid': get_xray_dict(anno_df, data_dir, cfg, valid_image_ids))
MetadataCatalog.get("xray_valid").set(thing_classes=list(class_name_dict.values()))
else:
train_image_ids = anno_df['image_id'].values
DatasetCatalog.register("xray_train", lambda d='train': get_xray_dict(anno_df, data_dir, cfg, train_image_ids))
MetadataCatalog.get("xray_train").set(thing_classes=list(class_name_dict.values()))
DatasetCatalog.register("xray_test", lambda d='test': get_test_xray_dict(data_dir))
MetadataCatalog.get("xray_test").set(thing_classes=list(class_name_dict.values()))
# Config --------------------------------------------------
detectron2_cfg = get_cfg()
detectron2_cfg.aug_kwargs = CN(rep_aug_kwargs)
detectron2_cfg.merge_from_file(model_zoo.get_config_file(backbone))
detectron2_cfg.DATASETS.TRAIN = ("xray_train",)
if 'valid' in cfg.data.split_method:
detectron2_cfg.DATASETS.TEST = ("xray_valid",)
detectron2_cfg.TEST.EVAL_PERIOD = cfg.train.max_iter // 10
else:
detectron2_cfg.DATASETS.TEST = ()
detectron2_cfg.INPUT.MIN_SIZE_TRAIN = (img_size,)
detectron2_cfg.INPUT.MAX_SIZE_TRAIN = img_size
detectron2_cfg.DATALOADER.NUM_WORKERS = cfg.train.num_workers
detectron2_cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(backbone)
detectron2_cfg.SOLVER.IMS_PER_BATCH = cfg.train.ims_per_batch
detectron2_cfg.SOLVER.BASE_LR = cfg.train.lr
detectron2_cfg.SOLVER.MAX_ITER = cfg.train.max_iter
detectron2_cfg.SOLVER.LR_SCHEDULER_NAME = "WarmupCosineLR"
detectron2_cfg.SOLVER.WARMUP_ITERS = 2000
detectron2_cfg.SOLVER.CHECKPOINT_PERIOD = 200000
detectron2_cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = cfg.train.batch_size_per_image
detectron2_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 15 if use_class14 else 14
detectron2_cfg.OUTPUT_DIR = output_dir
detectron2_cfg.SEED = cfg.data.seed
detectron2_cfg.PIXEL_MEAN = [103.530, 116.280, 123.675]
detectron2_cfg.PIXEL_STD = [1.0, 1.0, 1.0]
# Train --------------------------------------------------
os.makedirs(detectron2_cfg.OUTPUT_DIR, exist_ok=True)
# trainer = DefaultTrainer(detectron2_cfg)
trainer = MyTrainer(detectron2_cfg)
trainer.resume_or_load(resume=True)
trainer.train()
# Rename Last Weight
renamed_model = f"{backbone.split('.')[0].replace('/', '-')}.pth"
os.rename(os.path.join(cfg.data.output_dir, 'model_final.pth'),
os.path.join(cfg.data.output_dir, renamed_model))
# Logging
for model_path in glob.glob(os.path.join(cfg.data.output_dir, '*.pth')):
wandb.save(model_path)
# Inference Setting ------------------------------------------------------
detectron2_cfg = get_cfg()
detectron2_cfg.merge_from_file(model_zoo.get_config_file(backbone))
detectron2_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 15 if use_class14 else 14
detectron2_cfg.MODEL.WEIGHTS = os.path.join(output_dir, renamed_model) # path to the model we just trained
detectron2_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = cfg.data.score_th # set a custom testing threshold
predictor = DefaultPredictor(detectron2_cfg)
dataset_dicts = get_test_xray_dict(data_dir)
# Visualize ------------------------------------------------------
target_image_ids = ['9a5094b2563a1ef3ff50dc5c7ff71345',
'22b8e616a61bbc4caaed0cf23b7159df',
'001d127bad87592efe45a5c7678f8b8d',
'008b3176a7248a0a189b5731ac8d2e95']
for th in [0, 0.2, 0.5, 0.7]:
visualize(target_image_ids, data_dir, output_dir, predictor, score_th=th)
# Metrics
if os.path.exists(os.path.join(output_dir, 'metrics.json')):
metrics_df = pd.read_json(os.path.join(output_dir, 'metrics.json'), orient="records", lines=True)
mdf = metrics_df.sort_values("iteration")
mdf3 = mdf[~mdf["bbox/AP75"].isna()].reset_index(drop=True)
for i in range(len(mdf3)):
row = mdf3.iloc[i]
wandb.log({'AP40': row["bbox/AP75"] / 100.})
best_score = mdf3["bbox/AP75"].max() / 100.
wandb.log({'Best-AP40-Score': best_score})
# Inference ------------------------------------------------------
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
sub = get_submission(dataset_dicts, cfg, predictor, device)
now = datetime.datetime.now() + datetime.timedelta(hours=9)
now = now.strftime("%Y%m%d-%H%M%S")
filename = f'submission_{now}.csv'
sub.to_csv(os.path.join('./submission', filename), index=False)
wandb.save(os.path.join('./submission', filename))
time.sleep(30)
wandb.finish()
DatasetCatalog.clear()
def unregister():
DatasetCatalog.clear()
MetadataCatalog._NAME_TO_META = {}
def retrain_detector(settings):
"""
settings: properties to be used in the retraining process
Splits the COCO-formatted data located in annotation_path, then trains and
evaluates a Detectron2 model from scratch. The resulting model is saved in
the model_path/ folder.
Returns an object mapping different AP (average precision) metrics to the
model's scores.
"""
if len(settings) == 0:
settings["trainSplit"] = 0.7
settings["learningRate"] = 0.005
settings["maxIters"] = 100
base_path = "annotation_data/"
coco_path = os.path.join(base_path, "coco")
output_path = os.path.join(base_path, "output")
annotation_path = os.path.join(coco_path, "coco_results.json")
train_path = os.path.join(coco_path, "train.json")
test_path = os.path.join(coco_path, "test.json")
# 1) Split coco json file into train and test using cocosplit code
# Adapted from https://github.com/akarazniewicz/cocosplit/blob/master/cocosplit.py
with open(annotation_path, "rt", encoding="UTF-8") as annotations_file:
# Extract info from json
coco = json.load(annotations_file)
info = coco["info"]
licenses = coco["licenses"]
images = coco["images"]
annotations = coco["annotations"]
categories = coco["categories"]
# Remove images without annotations
images_with_annotations = set(
map(lambda a: int(a["image_id"]), annotations))
images = list(
filter(lambda i: i["id"] in images_with_annotations, images))
# Split images and annotations
x_images, y_images = train_test_split(
images, train_size=settings["trainSplit"])
x_ids = list(map(lambda i: int(i["id"]), x_images))
x_annots = list(
filter(lambda a: int(a["image_id"]) in x_ids, annotations))
y_ids = list(map(lambda i: int(i["id"]), y_images))
y_annots = list(
filter(lambda a: int(a["image_id"]) in y_ids, annotations))
# Save to file
def save_coco(file, info, licenses, images, annotations, categories):
with open(file, 'wt', encoding="UTF-8") as coco:
json.dump(
{
"info": info,
"licenses": licenses,
"images": images,
"annotations": annotations,
"categories": categories
},
coco,
indent=2,
sort_keys=True)
save_coco(train_path, info, licenses, x_images, x_annots, categories)
save_coco(test_path, info, licenses, y_images, y_annots, categories)
# 2) Use train/test files to retrain detector
dataset_name = "annotation_coco"
image_dir = base_path + "rgb/"
train_data = dataset_name + "_train"
test_data = dataset_name + "_test"
DatasetCatalog.clear()
MetadataCatalog.clear()
register_coco_instances(train_data, {}, train_path, image_dir)
register_coco_instances(test_data, {}, test_path, image_dir)
MetadataCatalog.get(train_data)
coco_yaml = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(coco_yaml))
cfg.DATASETS.TRAIN = (train_data, )
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(categories)
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
coco_yaml) # Let training initialize from model zoo
cfg.OUTPUT_DIR = output_path
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = settings["learningRate"] # Make sure LR is good
cfg.SOLVER.MAX_ITER = settings[
"maxIters"] # 300 is good for small datasets
# Train
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
# Move model to most recent model folder
model_dir = os.path.join(base_path, "model")
model_names = os.listdir(model_dir)
# Get highest x for model/vx
model_dirs = list(
filter(lambda n: os.path.isdir(os.path.join(model_dir, n)),
model_names))
model_nums = list(map(lambda x: int(x.split("v")[1]), model_dirs))
last_model_num = max(model_nums)
# Add model to new folder
model_path = os.path.join(model_dir, "v" + str(last_model_num))
new_model_path = os.path.join(model_path, "model_999.pth")
old_model_path = os.path.join(output_path, "model_final.pth")
os.replace(old_model_path, new_model_path)
# Evaluate
evaluator = COCOEvaluator(test_data, ("bbox", "segm"),
False,
output_dir="../../annotation_data/output/")
val_loader = build_detection_test_loader(cfg, test_data)
inference = inference_on_dataset(trainer.model, val_loader, evaluator)
# inference keys: bbox, semg
# bbox and segm keys: AP, AP50, AP75, APs, APm, AP1, AP-category1, ...
inference_json = json.loads(json.dumps(inference).replace("NaN", "null"))
return inference_json
def register_dataset(al_cfg, RUN_DIR, ALI, STRATEGY):
DatasetCatalog.clear()
TRAIN_IMAGE_IDS, VAL_IMAGE_IDS, FIXED_TEST_IMAGE_IDS, TEST_IMAGE_IDS = get_ids(
al_cfg, RUN_DIR, ALI, STRATEGY)
logger.info(
"Set Sizes\tSeed Set: {}\tVal Set: {}\tUnlabeled Set: {}".format(
len(TRAIN_IMAGE_IDS), len(VAL_IMAGE_IDS), len(TEST_IMAGE_IDS)))
IMG_IDS_DIC = {
'train': TRAIN_IMAGE_IDS,
'val': VAL_IMAGE_IDS,
'fixed': FIXED_TEST_IMAGE_IDS,
'test': TEST_IMAGE_IDS
}
if al_cfg.DATASET_NAME == "apollo":
apollo_classes = []
with open(os.path.join(al_cfg.DATASET_DIR,
'apollo_classes_things.txt')) as f:
for line in f.readlines():
apollo_classes.append(line.strip())
DatasetCatalog.register(
al_cfg.DATASETS.TRAIN,
lambda d='train': get_apollo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'RGB', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.VAL,
lambda d='val': get_apollo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'RGB', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.FIXED,
lambda d='fixed': get_apollo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'RGB', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.TEST,
lambda d='test': get_apollo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'RGB', d), IMG_IDS_DIC[d]))
MetadataCatalog.get(
al_cfg.DATASETS.TRAIN).set(thing_classes=apollo_classes)
MetadataCatalog.get(
al_cfg.DATASETS.VAL).set(thing_classes=apollo_classes)
MetadataCatalog.get(
al_cfg.DATASETS.FIXED).set(thing_classes=apollo_classes)
MetadataCatalog.get(
al_cfg.DATASETS.TEST).set(thing_classes=apollo_classes)
save_instances_per_class(
al_cfg, RUN_DIR, ALI, STRATEGY,
get_apollo_dicts(os.path.join(al_cfg.DATASET_DIR, 'RGB', 'train'),
TRAIN_IMAGE_IDS), apollo_classes)
elif al_cfg.DATASET_NAME == "yymnist":
yymnist_classes = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
DatasetCatalog.register(
al_cfg.DATASETS.TRAIN,
lambda d='train': get_yymnist_dicts(
os.path.join(al_cfg.DATASET_DIR, 'Images', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.VAL,
lambda d='val': get_yymnist_dicts(
os.path.join(al_cfg.DATASET_DIR, 'Images', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.FIXED,
lambda d='fixed': get_yymnist_dicts(
os.path.join(al_cfg.DATASET_DIR, 'Images', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.TEST,
lambda d='test': get_yymnist_dicts(
os.path.join(al_cfg.DATASET_DIR, 'Images', d), IMG_IDS_DIC[d]))
MetadataCatalog.get(
al_cfg.DATASETS.TRAIN).set(thing_classes=yymnist_classes)
MetadataCatalog.get(
al_cfg.DATASETS.VAL).set(thing_classes=yymnist_classes)
MetadataCatalog.get(
al_cfg.DATASETS.FIXED).set(thing_classes=yymnist_classes)
MetadataCatalog.get(
al_cfg.DATASETS.TEST).set(thing_classes=yymnist_classes)
save_instances_per_class(
al_cfg, RUN_DIR, ALI, STRATEGY,
get_yymnist_dicts(
os.path.join(al_cfg.DATASET_DIR, 'Images', 'train'),
TRAIN_IMAGE_IDS), yymnist_classes)
elif al_cfg.DATASET_NAME == "waymo":
waymo_classes = []
with open(os.path.join(al_cfg.DATASET_DIR,
'waymo_classes_things.txt')) as f:
for line in f.readlines():
waymo_classes.append(line.strip())
DatasetCatalog.register(
al_cfg.DATASETS.TRAIN,
lambda d='train': get_waymo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'RGB', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.VAL,
lambda d='val': get_waymo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'RGB', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.FIXED,
lambda d='fixed': get_waymo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'RGB', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.TEST,
lambda d='test': get_waymo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'RGB', d), IMG_IDS_DIC[d]))
MetadataCatalog.get(
al_cfg.DATASETS.TRAIN).set(thing_classes=waymo_classes)
MetadataCatalog.get(
al_cfg.DATASETS.VAL).set(thing_classes=waymo_classes)
MetadataCatalog.get(
al_cfg.DATASETS.FIXED).set(thing_classes=waymo_classes)
MetadataCatalog.get(
al_cfg.DATASETS.TEST).set(thing_classes=waymo_classes)
save_instances_per_class(
al_cfg, RUN_DIR, ALI, STRATEGY,
get_waymo_dicts(os.path.join(al_cfg.DATASET_DIR, 'RGB', 'train'),
TRAIN_IMAGE_IDS), waymo_classes)
elif al_cfg.DATASET_NAME == "project_waymo":
project_waymo_classes = ['vehicle', 'person', 'cyclist', 'sign']
DatasetCatalog.register(
al_cfg.DATASETS.TRAIN,
lambda d='train': get_project_waymo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'images', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.VAL,
lambda d='val': get_project_waymo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'images', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.FIXED,
lambda d='fixed': get_project_waymo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'images', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.TEST,
lambda d='test': get_project_waymo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'images', d), IMG_IDS_DIC[d]))
MetadataCatalog.get(
al_cfg.DATASETS.TRAIN).set(thing_classes=project_waymo_classes)
MetadataCatalog.get(
al_cfg.DATASETS.VAL).set(thing_classes=project_waymo_classes)
MetadataCatalog.get(
al_cfg.DATASETS.FIXED).set(thing_classes=project_waymo_classes)
MetadataCatalog.get(
al_cfg.DATASETS.TEST).set(thing_classes=project_waymo_classes)
save_instances_per_class(
al_cfg, RUN_DIR, ALI, STRATEGY,
get_project_waymo_dicts(
os.path.join(al_cfg.DATASET_DIR, 'images', 'train'),
TRAIN_IMAGE_IDS), project_waymo_classes)
elif al_cfg.DATASET_NAME == "naplab":
DatasetCatalog.register(
al_cfg.DATASETS.TRAIN,
lambda d='train': get_naplab_dicts(
os.path.join(al_cfg.DATASET_DIR, 'images', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.VAL,
lambda d='val': get_naplab_dicts(
os.path.join(al_cfg.DATASET_DIR, 'images', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.FIXED,
lambda d='fixed': get_naplab_dicts(
os.path.join(al_cfg.DATASET_DIR, 'images', d), IMG_IDS_DIC[d]))
DatasetCatalog.register(
al_cfg.DATASETS.TEST,
lambda d='test': get_naplab_dicts(
os.path.join(al_cfg.DATASET_DIR, 'images', d), IMG_IDS_DIC[d]))
MetadataCatalog.get(al_cfg.DATASETS.TRAIN).set(thing_classes=[])
MetadataCatalog.get(al_cfg.DATASETS.VAL).set(thing_classes=[])
MetadataCatalog.get(al_cfg.DATASETS.FIXED).set(thing_classes=[])
MetadataCatalog.get(al_cfg.DATASETS.TEST).set(thing_classes=[])
# save_instances_per_class(al_cfg, RUN_DIR, ALI, STRATEGY, get_naplab_dicts(os.path.join(al_cfg.DATASET_DIR, 'images', 'train'), TRAIN_IMAGE_IDS), [])
else:
logger.error(
"Dataset not registered! Code not implemented for dataset: {}".
format(al_cfg.DATASET_NAME))
raise NotImplementedError
