Add models implementation and training outputs

faster_rcnn_R_101_FPN_3x/eval.py

@@ -0,0 +1,44 @@
+from detectron2.data import DatasetCatalog, DatasetMapper
+from detectron2.engine import DefaultPredictor
+from detectron2.evaluation import COCOEvaluator, inference_on_dataset
+from detectron2.data import build_detection_test_loader
+from utils import (
+    build_config,
+    register_publaynet_datasets,
+)
+
+# Register datasets
+dataset_train_name, dataset_test_name = register_publaynet_datasets()
+
+# Model parameters
+model_zoo_config_name = "COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"
+trained_model_output_dir = "/home/alex/Dev/deep_learning/detectron2-publaynet/faster_rcnn_R_101_FPN_3x/training_output"
+prediction_score_threshold = 0.7
+base_lr = 0.001
+max_iter = 500
+batch_size = 128
+
+# Detectron config
+cfg = build_config(
+    model_zoo_config_name,
+    dataset_train_name,
+    dataset_test_name,
+    trained_model_output_dir,
+    prediction_score_threshold,
+    base_lr,
+    max_iter,
+    batch_size,
+)
+
+# Detectron predictor
+predictor = DefaultPredictor(cfg)
+
+# Load test dataset, and evaluate over it
+dataset = DatasetCatalog.get(dataset_test_name)
+dataset_mapper = DatasetMapper(cfg, is_train=False)
+data_loader = build_detection_test_loader(dataset=dataset, mapper=dataset_mapper)
+inference_on_dataset(
+    model=predictor.model,
+    data_loader=data_loader,
+    evaluator=COCOEvaluator(dataset_test_name),
+)

faster_rcnn_R_101_FPN_3x/evaluation.txt

@@ -0,0 +1,33 @@
+[09/20 16:54:58 d2.evaluation.evaluator]: Total inference time: 0:19:19.760338 (0.103182 s / iter per device, on 1 devices)
+[09/20 16:54:58 d2.evaluation.evaluator]: Total inference pure compute time: 0:15:54 (0.084938 s / iter per device, on 1 devices)
+[09/20 16:54:58 d2.evaluation.coco_evaluation]: Preparing results for COCO format ...
+[09/20 16:54:58 d2.evaluation.coco_evaluation]: Evaluating predictions with unofficial COCO API...
+Loading and preparing results...
+DONE (t=0.40s)
+creating index...
+index created!
+[09/20 16:54:59 d2.evaluation.fast_eval_api]: Evaluate annotation type *bbox*
+[09/20 16:55:01 d2.evaluation.fast_eval_api]: COCOeval_opt.evaluate() finished in 2.51 seconds.
+[09/20 16:55:01 d2.evaluation.fast_eval_api]: Accumulating evaluation results...
+[09/20 16:55:02 d2.evaluation.fast_eval_api]: COCOeval_opt.accumulate() finished in 0.39 seconds.
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.843
+ Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.940
+ Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.912
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.251
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.602
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.894
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.503
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.871
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.880
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.270
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.650
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.927
+[09/20 16:55:02 d2.evaluation.coco_evaluation]: Evaluation results for bbox:
+|   AP   |  AP50  |  AP75  |  APs   |  APm   |  APl   |
+|:------:|:------:|:------:|:------:|:------:|:------:|
+| 84.295 | 94.048 | 91.199 | 25.061 | 60.217 | 89.431 |
+[09/20 16:55:02 d2.evaluation.coco_evaluation]: Per-category bbox AP:
+| category   | AP     | category   | AP     | category   | AP     |
+|:-----------|:-------|:-----------|:-------|:-----------|:-------|
+| text       | 89.249 | title      | 76.824 | list       | 77.612 |
+| table      | 92.250 | figure     | 85.539 |            |        |

faster_rcnn_R_101_FPN_3x/test.py

@@ -0,0 +1,29 @@
+from detectron2.engine import DefaultPredictor
+from utils import build_config, register_publaynet_datasets, visual_test
+
+# Register datasets
+dataset_train_name, dataset_test_name = register_publaynet_datasets()
+
+# Model parameters
+model_zoo_config_name = "COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"
+trained_model_output_dir = "/home/alex/Dev/deep_learning/detectron2-publaynet/faster_rcnn_R_101_FPN_3x/training_output"
+prediction_score_threshold = 0.7
+base_lr = 0.001
+max_iter = 50
+batch_size = 128
+
+# Detectron config
+cfg = build_config(
+    model_zoo_config_name,
+    dataset_train_name,
+    dataset_test_name,
+    trained_model_output_dir,
+    prediction_score_threshold,
+    base_lr,
+    max_iter,
+    batch_size,
+)
+
+# Detectron predictor
+predictor = DefaultPredictor(cfg)
+visual_test(cfg, predictor)

faster_rcnn_R_101_FPN_3x/train.py

@@ -0,0 +1,30 @@
+from detectron2.engine import DefaultTrainer
+from utils import build_config, register_publaynet_datasets
+
+# Register datasets
+dataset_train_name, dataset_test_name = register_publaynet_datasets()
+
+# Model parameters
+model_zoo_config_name = "COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"
+trained_model_output_dir = "/home/alex/Dev/deep_learning/detectron2-publaynet/faster_rcnn_R_101_FPN_3x/training_output"
+prediction_score_threshold = 0.7
+base_lr = 0.00001
+max_iter = 8000
+batch_size = 128
+
+# Detectron config
+cfg = build_config(
+    model_zoo_config_name,
+    dataset_train_name,
+    dataset_test_name,
+    trained_model_output_dir,
+    prediction_score_threshold,
+    base_lr,
+    max_iter,
+    batch_size,
+)
+
+# Detectron Trainer
+trainer = DefaultTrainer(cfg)
+trainer.resume_or_load(resume=False)
+trainer.train()

faster_rcnn_R_101_FPN_3x/training_output/training_output/events.out.tfevents.1726006373.Legion.31775.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31f639edc57729dc50b05d14ce2101eda3bfb966037a64542bbef8d394cd7f55
+size 91186

faster_rcnn_R_101_FPN_3x/training_output/training_output/events.out.tfevents.1726406883.Legion.19405.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3b2bd51b5d8529ecff5fb5cf2999aa052a1c4958972c0f3728528484813dc02
+size 364786

faster_rcnn_R_101_FPN_3x/training_output/training_output/last_checkpoint

@@ -0,0 +1 @@
+model_final.pth

faster_rcnn_R_101_FPN_3x/training_output/training_output/model_0004999.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5854879a3f88d2ccc9ac871a4dee4a79ad3d0d7e73dae646d031ec1d2eb4901e
+size 482234736

faster_rcnn_R_101_FPN_3x/training_output/training_output/model_final.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e58306b214ec6c7fef0e5daa14eac84a546d2313322c2dcd4946a6338f02589
+size 482234736

faster_rcnn_R_101_FPN_3x/utils.py

@@ -0,0 +1,114 @@
+import os
+from detectron2.utils.visualizer import Visualizer
+import cv2
+import numpy
+from PIL.Image import Image
+from PIL import Image as image_main
+from detectron2.utils.logger import setup_logger
+
+setup_logger()
+from detectron2 import model_zoo
+from detectron2.config import get_cfg, CfgNode
+from detectron2.data import MetadataCatalog
+from detectron2.data.datasets.register_coco import register_coco_instances
+
+
+def open_image_pil(image_path: str) -> Image:
+    return image_main.open(image_path)
+
+
+def convert_pil_to_cv(pil_image: Image):
+    if pil_image.mode != "RGB":
+        pil_image = pil_image.convert("RGB")
+    return cv2.cvtColor(numpy.array(pil_image), cv2.COLOR_RGB2BGR)
+
+
+def register_publaynet_datasets() -> (str, str):
+    dataset_train_name = "publaynet_dataset_train"
+    dataset_test_name = "publaynet_dataset_test"
+    class_labels = ["text", "title", "list", "table", "figure"]
+    register_coco_instances(
+        dataset_train_name,
+        {},
+        "/home/alex/Datasets/PubLayNet/publaynet/train.json",
+        "/home/alex/Datasets/PubLayNet/publaynet/train",
+    )
+    register_coco_instances(
+        dataset_test_name,
+        {},
+        "/home/alex/Datasets/PubLayNet/publaynet/val.json",
+        "/home/alex/Datasets/PubLayNet/publaynet/val",
+    )
+
+    # Make sure the datasets got registered
+    metadata_train = MetadataCatalog.get(dataset_train_name)
+    metadata_test = MetadataCatalog.get(dataset_test_name)
+    print(metadata_train)
+    print(metadata_test)
+
+    # Set labels
+    MetadataCatalog.get(dataset_train_name).thing_classes = class_labels
+    MetadataCatalog.get(dataset_test_name).thing_classes = class_labels
+
+    return dataset_train_name, dataset_test_name
+
+
+def build_config(
+    model_zoo_config_name: str,
+    dataset_train_name: str,
+    dataset_test_name: str,
+    trained_model_output_dir: str,
+    prediction_score_threshold: float,
+    base_lr: float,
+    max_iter: int,
+    batch_size: int,
+) -> CfgNode:
+    trained_model_weights_path = trained_model_output_dir + "/model_final.pth"
+
+    cfg = get_cfg()
+    cfg.merge_from_file(model_zoo.get_config_file(model_zoo_config_name))
+    cfg.DATASETS.TRAIN = (dataset_train_name,)
+    cfg.DATASETS.TEST = (dataset_test_name,)
+    cfg.OUTPUT_DIR = trained_model_output_dir
+    cfg.DATALOADER.NUM_WORKERS = 8
+    if os.path.exists(trained_model_weights_path):
+        cfg.MODEL.WEIGHTS = trained_model_weights_path
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = prediction_score_threshold
+    cfg.SOLVER.IMS_PER_BATCH = 4
+    cfg.SOLVER.BASE_LR = base_lr
+    cfg.SOLVER.MAX_ITER = max_iter
+    cfg.SOLVER.STEPS = []
+    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = batch_size
+    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 5
+    cfg.TEST.DETECTIONS_PER_IMAGE = 100
+    # cfg.INPUT.MIN_SIZE_TRAIN = (640, 672, 704, 736, 768, 800)
+    # cfg.INPUT.MIN_SIZE_TRAIN = (600, 632, 664, 696, 728, 760)
+    cfg.INPUT.MIN_SIZE_TRAIN = (580, 612, 644, 676, 708, 740)
+    return cfg
+
+
+def visualize_outputs(cfg, image_cv, outputs, name):
+    v = Visualizer(
+        image_cv[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2
+    )
+    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+    output_image_cv = out.get_image()[:, :, ::-1]
+    cv2.imwrite(name, output_image_cv)
+
+
+def visual_test(cfg: CfgNode, predictor: DefaultPredictor):
+    image_paths = [
+        "/home/alex/Datasets/PubLayNet/publaynet/train/PMC1500815_00002.jpg",
+        "/home/alex/Datasets/PubLayNet/publaynet/train/PMC3162874_00002.jpg",
+        "/home/alex/Datasets/PubLayNet/publaynet/train/PMC4203354_00000.jpg",
+        "/home/alex/Datasets/PubLayNet/publaynet/val/PMC1247188_00003.jpg",
+        "/home/alex/Datasets/PubLayNet/publaynet/val/PMC2829689_00004.jpg",
+        "/home/alex/Datasets/PubLayNet/publaynet/val/PMC4520132_00000.jpg",
+    ]
+
+    for i, image_path in enumerate(image_paths):
+        print("Testing on " + image_path)
+        image_pil = open_image_pil(image_path)
+        image_cv = convert_pil_to_cv(image_pil)
+        outputs = predictor(image_cv)
+        visualize_outputs(cfg, image_cv, outputs, f"image_{i}.jpg")

mask_rcnn_R_101_FPN_3x/eval.py

@@ -0,0 +1,44 @@
+from detectron2.data import DatasetCatalog, DatasetMapper
+from detectron2.engine import DefaultPredictor
+from detectron2.evaluation import COCOEvaluator, inference_on_dataset
+from detectron2.data import build_detection_test_loader
+from utils import (
+    build_config,
+    register_publaynet_datasets,
+)
+
+# Register datasets
+dataset_train_name, dataset_test_name = register_publaynet_datasets()
+
+# Model parameters
+model_zoo_config_name = "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"
+trained_model_output_dir = "/home/alex/Dev/deep_learning/detectron2-publaynet/mask_rcnn_R_101_FPN_3x/training_output"
+prediction_score_threshold = 0.7
+base_lr = 0.0001
+max_iter = 1000
+batch_size = 128
+
+# Detectron config
+cfg = build_config(
+    model_zoo_config_name,
+    dataset_train_name,
+    dataset_test_name,
+    trained_model_output_dir,
+    prediction_score_threshold,
+    base_lr,
+    max_iter,
+    batch_size,
+)
+
+# Detectron predictor
+predictor = DefaultPredictor(cfg)
+
+# Load test dataset, and evaluate over it
+dataset = DatasetCatalog.get(dataset_test_name)
+dataset_mapper = DatasetMapper(cfg, is_train=False)
+data_loader = build_detection_test_loader(dataset=dataset, mapper=dataset_mapper)
+inference_on_dataset(
+    model=predictor.model,
+    data_loader=data_loader,
+    evaluator=COCOEvaluator(dataset_test_name),
+)

mask_rcnn_R_101_FPN_3x/evaluation.txt

@@ -0,0 +1,62 @@
+[09/26 01:03:24 d2.evaluation.evaluator]: Total inference time: 0:20:57.157306 (0.111847 s / iter per device, on 1 devices)
+[09/26 01:03:24 d2.evaluation.evaluator]: Total inference pure compute time: 0:16:27 (0.087851 s / iter per device, on 1 devices)
+[09/26 01:03:24 d2.evaluation.coco_evaluation]: Preparing results for COCO format ...
+[09/26 01:03:24 d2.evaluation.coco_evaluation]: Evaluating predictions with unofficial COCO API...
+Loading and preparing results...
+DONE (t=0.06s)
+creating index...
+index created!
+[09/26 01:03:24 d2.evaluation.fast_eval_api]: Evaluate annotation type *bbox*
+[09/26 01:03:27 d2.evaluation.fast_eval_api]: COCOeval_opt.evaluate() finished in 2.66 seconds.
+[09/26 01:03:27 d2.evaluation.fast_eval_api]: Accumulating evaluation results...
+[09/26 01:03:27 d2.evaluation.fast_eval_api]: COCOeval_opt.accumulate() finished in 0.41 seconds.
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.867
+ Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.948
+ Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.923
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.290
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.615
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.916
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.514
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.889
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.898
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.313
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.658
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.942
+[09/26 01:03:27 d2.evaluation.coco_evaluation]: Evaluation results for bbox:
+|   AP   |  AP50  |  AP75  |  APs   |  APm   |  APl   |
+|:------:|:------:|:------:|:------:|:------:|:------:|
+| 86.690 | 94.839 | 92.308 | 29.010 | 61.459 | 91.559 |
+[09/26 01:03:27 d2.evaluation.coco_evaluation]: Per-category bbox AP:
+| category   | AP     | category   | AP     | category   | AP     |
+|:-----------|:-------|:-----------|:-------|:-----------|:-------|
+| text       | 89.822 | title      | 79.101 | list       | 80.716 |
+| table      | 94.215 | figure     | 89.594 |            |        |
+Loading and preparing results...
+DONE (t=1.41s)
+creating index...
+index created!
+[09/26 01:03:30 d2.evaluation.fast_eval_api]: Evaluate annotation type *segm*
+[09/26 01:03:39 d2.evaluation.fast_eval_api]: COCOeval_opt.evaluate() finished in 8.26 seconds.
+[09/26 01:03:39 d2.evaluation.fast_eval_api]: Accumulating evaluation results...
+[09/26 01:03:39 d2.evaluation.fast_eval_api]: COCOeval_opt.accumulate() finished in 0.39 seconds.
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.821
+ Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.947
+ Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.898
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.280
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.569
+ Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.862
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.489
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.854
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.862
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.311
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.620
+ Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.898
+[09/26 01:03:39 d2.evaluation.coco_evaluation]: Evaluation results for segm:
+|   AP   |  AP50  |  AP75  |  APs   |  APm   |  APl   |
+|:------:|:------:|:------:|:------:|:------:|:------:|
+| 82.105 | 94.654 | 89.840 | 28.016 | 56.863 | 86.208 |
+[09/26 01:03:39 d2.evaluation.coco_evaluation]: Per-category segm AP:
+| category   | AP     | category   | AP     | category   | AP     |
+|:-----------|:-------|:-----------|:-------|:-----------|:-------|
+| text       | 88.786 | title      | 76.630 | list       | 62.243 |
+| table      | 93.647 | figure     | 89.217 |            |        |

mask_rcnn_R_101_FPN_3x/test.py

@@ -0,0 +1,29 @@
+from detectron2.engine import DefaultPredictor
+from utils import build_config, register_publaynet_datasets, visual_test
+
+# Register datasets
+dataset_train_name, dataset_test_name = register_publaynet_datasets()
+
+# Model parameters
+model_zoo_config_name = "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"
+trained_model_output_dir = "/home/alex/Dev/deep_learning/detectron2-publaynet/mask_rcnn_R_101_FPN_3x/training_output"
+prediction_score_threshold = 0.7
+base_lr = 0
+max_iter = 0
+batch_size = 0
+
+# Detectron config
+cfg = build_config(
+    model_zoo_config_name,
+    dataset_train_name,
+    dataset_test_name,
+    trained_model_output_dir,
+    prediction_score_threshold,
+    base_lr,
+    max_iter,
+    batch_size,
+)
+
+# Detectron predictor
+predictor = DefaultPredictor(cfg)
+visual_test(cfg, predictor)

mask_rcnn_R_101_FPN_3x/train.py

@@ -0,0 +1,30 @@
+from detectron2.engine import DefaultTrainer
+from utils import build_config, register_publaynet_datasets
+
+# Register datasets
+dataset_train_name, dataset_test_name = register_publaynet_datasets()
+
+# Model parameters
+model_zoo_config_name = "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"
+trained_model_output_dir = "/home/alex/Dev/deep_learning/detectron2-publaynet/mask_rcnn_R_101_FPN_3x/training_output"
+prediction_score_threshold = 0.7
+base_lr = 0.00001
+max_iter = 15000
+batch_size = 128
+
+# Detectron config
+cfg = build_config(
+    model_zoo_config_name,
+    dataset_train_name,
+    dataset_test_name,
+    trained_model_output_dir,
+    prediction_score_threshold,
+    base_lr,
+    max_iter,
+    batch_size,
+)
+
+# Detectron Trainer
+trainer = DefaultTrainer(cfg)
+trainer.resume_or_load(resume=False)
+trainer.train()

mask_rcnn_R_101_FPN_3x/training_output/training_output/events.out.tfevents.1725897998.Legion.47938.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f13688b1d5c48925ef86b9abfbc8caddf0b77929a4e635754d1f6be4154023d5
+size 57112

mask_rcnn_R_101_FPN_3x/training_output/training_output/events.out.tfevents.1726172796.Legion.3380.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:260f4db66bc80b5c0a0cf0bc60d3e1837bcee0d0653c76f2f485c578f3741d29
+size 171412

mask_rcnn_R_101_FPN_3x/training_output/training_output/events.out.tfevents.1726421934.Legion.80102.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ddd5e4f8e0de29661376bbe5124c4e40cca3e6158be0c2dc75a5af118deb8a6
+size 20536

mask_rcnn_R_101_FPN_3x/training_output/training_output/events.out.tfevents.1726422787.Legion.1493.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0db25f7e76307a7515fd271c972f11a6a8fe8bd9e56232e0811c835850f38cb0
+size 20536

mask_rcnn_R_101_FPN_3x/training_output/training_output/events.out.tfevents.1726423456.Legion.4616.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32f5311e5c29787029c04da765c5bf0e735c78cd6d0609116a07f9f15ae38c88
+size 17107

mask_rcnn_R_101_FPN_3x/training_output/training_output/events.out.tfevents.1726423883.Legion.6492.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92134b9c0422d1a7ad420c392d6c1659b867e8e0418011d3116e456495f6da6f
+size 857212

mask_rcnn_R_101_FPN_3x/training_output/training_output/last_checkpoint

@@ -0,0 +1 @@
+model_final.pth

mask_rcnn_R_101_FPN_3x/training_output/training_output/model_0004999.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d0951119dc4443fd4b3fa7b8a963251e6fd043bf1816ee403da5a649f71881b
+size 503235392

mask_rcnn_R_101_FPN_3x/training_output/training_output/model_0009999.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c84659e4c6e5844940acb6909032f7aee1947818f9d1e2098ceabaa15cc3c579
+size 503235392

mask_rcnn_R_101_FPN_3x/training_output/training_output/model_0014999.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df3036f36789631a59f532fe615f13d4ee1ba13aafd420f1c1cef8536d5cb1fc
+size 503235392

mask_rcnn_R_101_FPN_3x/training_output/training_output/model_final.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df3036f36789631a59f532fe615f13d4ee1ba13aafd420f1c1cef8536d5cb1fc
+size 503235392

mask_rcnn_R_101_FPN_3x/utils.py

@@ -0,0 +1,117 @@
+import os
+from typing import List
+from detectron2.utils.visualizer import Visualizer
+import cv2
+import numpy
+from PIL.Image import Image
+from PIL import Image as image_main
+import detectron2
+from detectron2.utils.logger import setup_logger
+
+setup_logger()
+from detectron2 import model_zoo
+from detectron2.config import get_cfg, CfgNode
+from detectron2.engine import DefaultTrainer, DefaultPredictor
+from detectron2.data import MetadataCatalog
+from detectron2.data.datasets.register_coco import register_coco_instances
+
+
+def open_image_pil(image_path: str) -> Image:
+    return image_main.open(image_path)
+
+
+def convert_pil_to_cv(pil_image: Image):
+    if pil_image.mode != "RGB":
+        pil_image = pil_image.convert("RGB")
+    return cv2.cvtColor(numpy.array(pil_image), cv2.COLOR_RGB2BGR)
+
+
+def register_publaynet_datasets() -> (str, str):
+    dataset_train_name = "publaynet_dataset_train"
+    dataset_test_name = "publaynet_dataset_test"
+    class_labels = ["text", "title", "list", "table", "figure"]
+    register_coco_instances(
+        dataset_train_name,
+        {},
+        "/home/alex/Datasets/PubLayNet/publaynet/train.json",
+        "/home/alex/Datasets/PubLayNet/publaynet/train",
+    )
+    register_coco_instances(
+        dataset_test_name,
+        {},
+        "/home/alex/Datasets/PubLayNet/publaynet/val.json",
+        "/home/alex/Datasets/PubLayNet/publaynet/val",
+    )
+
+    # Make sure the datasets got registered
+    metadata_train = MetadataCatalog.get(dataset_train_name)
+    metadata_test = MetadataCatalog.get(dataset_test_name)
+    print(metadata_train)
+    print(metadata_test)
+
+    # Set labels
+    MetadataCatalog.get(dataset_train_name).thing_classes = class_labels
+    MetadataCatalog.get(dataset_test_name).thing_classes = class_labels
+
+    return dataset_train_name, dataset_test_name
+
+
+def build_config(
+    model_zoo_config_name: str,
+    dataset_train_name: str,
+    dataset_test_name: str,
+    trained_model_output_dir: str,
+    prediction_score_threshold: float,
+    base_lr: float,
+    max_iter: int,
+    batch_size: int,
+) -> CfgNode:
+    trained_model_weights_path = trained_model_output_dir + "/model_final.pth"
+
+    cfg = get_cfg()
+    cfg.merge_from_file(model_zoo.get_config_file(model_zoo_config_name))
+    cfg.DATASETS.TRAIN = (dataset_train_name,)
+    cfg.DATASETS.TEST = (dataset_test_name,)
+    cfg.OUTPUT_DIR = trained_model_output_dir
+    cfg.DATALOADER.NUM_WORKERS = 8
+    if os.path.exists(trained_model_weights_path):
+        cfg.MODEL.WEIGHTS = trained_model_weights_path
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = prediction_score_threshold
+    cfg.SOLVER.IMS_PER_BATCH = 4
+    cfg.SOLVER.BASE_LR = base_lr
+    cfg.SOLVER.MAX_ITER = max_iter
+    cfg.SOLVER.STEPS = []
+    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = batch_size
+    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 5
+    cfg.TEST.DETECTIONS_PER_IMAGE = 100
+    # cfg.INPUT.MIN_SIZE_TRAIN = (640, 672, 704, 736, 768, 800)
+    # cfg.INPUT.MIN_SIZE_TRAIN = (600, 632, 664, 696, 728, 760)
+    cfg.INPUT.MIN_SIZE_TRAIN = (580, 612, 644, 676, 708, 740)
+    return cfg
+
+
+def visualize_outputs(cfg, image_cv, outputs, name):
+    v = Visualizer(
+        image_cv[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2
+    )
+    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+    output_image_cv = out.get_image()[:, :, ::-1]
+    cv2.imwrite(name, output_image_cv)
+
+
+def visual_test(cfg: CfgNode, predictor: DefaultPredictor):
+    image_paths = [
+        "/home/alex/Datasets/PubLayNet/publaynet/train/PMC1500815_00002.jpg",
+        "/home/alex/Datasets/PubLayNet/publaynet/train/PMC3162874_00002.jpg",
+        "/home/alex/Datasets/PubLayNet/publaynet/train/PMC4203354_00000.jpg",
+        "/home/alex/Datasets/PubLayNet/publaynet/val/PMC1247188_00003.jpg",
+        "/home/alex/Datasets/PubLayNet/publaynet/val/PMC2829689_00004.jpg",
+        "/home/alex/Datasets/PubLayNet/publaynet/val/PMC4520132_00000.jpg",
+    ]
+
+    for i, image_path in enumerate(image_paths):
+        print("Testing on " + image_path)
+        image_pil = open_image_pil(image_path)
+        image_cv = convert_pil_to_cv(image_pil)
+        outputs = predictor(image_cv)
+        visualize_outputs(cfg, image_cv, outputs, f"image_{i}.jpg")