Upload 5 files

detection/configs/Base-RCNN-FPN.yaml

@@ -0,0 +1,44 @@
+MODEL:
+  SEM_SEG_HEAD:
+    NUM_CLASSES: 15
+  META_ARCHITECTURE: "GeneralizedRCNN"
+  BACKBONE:
+    NAME: "build_resnet_fpn_backbone"
+  RESNETS:
+    OUT_FEATURES: ["res2", "res3", "res4", "res5"]
+  FPN:
+    IN_FEATURES: ["res2", "res3", "res4", "res5"]
+  ANCHOR_GENERATOR:
+    SIZES: [[32], [64], [128], [256], [512]]  # One size for each in feature map
+    ASPECT_RATIOS: [[0.5, 1.0, 2.0]]  # Three aspect ratios (same for all in feature maps)
+  RPN:
+    IN_FEATURES: ["p2", "p3", "p4", "p5", "p6"]
+    PRE_NMS_TOPK_TRAIN: 2000  # Per FPN level
+    PRE_NMS_TOPK_TEST: 1000  # Per FPN level
+    # Detectron1 uses 2000 proposals per-batch,
+    # (See "modeling/rpn/rpn_outputs.py" for details of this legacy issue)
+    # which is approximately 1000 proposals per-image since the default batch size for FPN is 2.
+    POST_NMS_TOPK_TRAIN: 1000
+    POST_NMS_TOPK_TEST: 1000
+  ROI_HEADS:
+    NAME: "StandardROIHeads"
+    IN_FEATURES: ["p2", "p3", "p4", "p5"]
+  ROI_BOX_HEAD:
+    NAME: "FastRCNNConvFCHead"
+    NUM_FC: 2
+    POOLER_RESOLUTION: 7
+  ROI_MASK_HEAD:
+    NAME: "MaskRCNNConvUpsampleHead"
+    NUM_CONV: 4
+    POOLER_RESOLUTION: 14
+DATASETS:
+  TRAIN: ("train_dora_ui",)
+  TEST: ("valid_dora_ui",)
+SOLVER:
+  IMS_PER_BATCH: 16
+  BASE_LR: 0.0005
+  STEPS: (60000, 80000)
+  MAX_ITER: 90000
+INPUT:
+  MIN_SIZE_TRAIN: (640, 672, 704, 736, 768, 800)
+VERSION: 2

detection/configs/faster_rcnn_R_50_FPN_1x.yaml

@@ -0,0 +1,10 @@
+_BASE_: "Base-RCNN-FPN.yaml"
+MODEL:
+  WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl"
+  MASK_ON: False
+  RESNETS:
+    DEPTH: 50
+SOLVER:
+  CHECKPOINT_PERIOD: 5000
+TEST:
+  EVAL_PERIOD: 1000

detection/data_util.py

@@ -0,0 +1,43 @@
+import json
+import numpy as np
+
+from icecream import ic, install
+install()
+ic.configureOutput(includeContext=True, contextAbsPath=True)
+
+def load_coco_json(json_path):
+    with open(json_path, 'r') as f:
+        data = json.load(f)
+    return data
+
+# ['info', 'licenses', 'images', 'annotations', 'categories']
+
+def split_train_val(data, val_ratio=0.1):
+    img_ids = [img['id'] for img in data['images']]
+    img_ids = np.array(img_ids)
+    np.random.shuffle(img_ids)
+    val_num = int(len(img_ids) * val_ratio)
+    val_ids = img_ids[:val_num]
+    train_ids = img_ids[val_num:]
+    train_data = {'info': data['info'], 'licenses': data['licenses'], 'images': [], 'annotations': [], 'categories': data['categories']}
+    val_data = {'info': data['info'], 'licenses': data['licenses'], 'images': [], 'annotations': [], 'categories': data['categories']}
+    for img in data['images']:
+        if img['id'] in train_ids:
+            train_data['images'].append(img)
+        else:
+            val_data['images'].append(img)
+    for ann in data['annotations']:
+        if ann['image_id'] in train_ids:
+            train_data['annotations'].append(ann)
+        else:
+            val_data['annotations'].append(ann)
+    return train_data, val_data
+
+data = load_coco_json('/root/autodl-tmp/dora_dataset/train/_annotations.coco.json')
+train_data, val_data = split_train_val(data)
+# save train_data and val_data
+with open('/root/autodl-tmp/dora_dataset/train.json', 'w') as f:
+    json.dump(train_data, f)
+with open('/root/autodl-tmp/dora_dataset/val.json', 'w') as f:
+    json.dump(val_data, f)
+

detection/requirements.txt

@@ -0,0 +1,134 @@
+absl-py==1.0.0
+antlr4-python3-runtime==4.9.3
+anyio==3.4.0
+appdirs==1.4.4
+argon2-cffi==21.1.0
+asttokens==2.2.1
+attrs==21.2.0
+Babel==2.9.1
+backcall==0.2.0
+black==21.4b2
+bleach==4.1.0
+brotlipy==0.7.0
+cachetools==4.2.4
+certifi==2021.5.30
+cffi @ file:///tmp/build/80754af9/cffi_1625807838443/work
+chardet @ file:///tmp/build/80754af9/chardet_1607706746162/work
+click==8.1.3
+cloudpickle==2.1.0
+colorama==0.4.5
+conda==4.10.3
+conda-package-handling @ file:///tmp/build/80754af9/conda-package-handling_1618262148928/work
+cryptography @ file:///tmp/build/80754af9/cryptography_1616769286105/work
+cycler==0.11.0
+debugpy==1.5.1
+decorator==5.1.0
+defusedxml==0.7.1
+detectron2==0.6+cu113
+entrypoints==0.3
+executing==1.2.0
+fonttools==4.28.2
+future==0.18.2
+fvcore==0.1.5.post20221122
+google-auth==2.3.3
+google-auth-oauthlib==0.4.6
+grpcio==1.42.0
+hydra-core==1.3.2
+icecream==2.1.3
+idna @ file:///home/linux1/recipes/ci/idna_1610986105248/work
+importlib-metadata==4.8.2
+importlib-resources==5.4.0
+iopath==0.1.9
+ipykernel==6.5.1
+ipython==7.29.0
+ipython-genutils==0.2.0
+ipywidgets==7.6.5
+jedi==0.18.1
+Jinja2==3.0.3
+json5==0.9.6
+jsonschema==4.2.1
+jupyter-client==7.1.0
+jupyter-core==4.9.1
+jupyter-server==1.12.0
+jupyterlab==3.2.4
+jupyterlab-language-pack-zh-CN==3.2.post2
+jupyterlab-pygments==0.1.2
+jupyterlab-server==2.8.2
+jupyterlab-widgets==1.0.2
+kiwisolver==1.3.2
+Markdown==3.3.6
+MarkupSafe==2.0.1
+matplotlib==3.5.0
+matplotlib-inline==0.1.3
+mistune==0.8.4
+mypy-extensions==1.0.0
+nbclassic==0.3.4
+nbclient==0.5.9
+nbconvert==6.3.0
+nbformat==5.1.3
+nest-asyncio==1.5.1
+notebook==6.4.6
+numpy==1.21.4
+oauthlib==3.1.1
+omegaconf==2.3.0
+opencv-python==4.7.0.72
+packaging==21.3
+pandocfilters==1.5.0
+parso==0.8.2
+pathspec==0.11.0
+pexpect==4.8.0
+pickleshare==0.7.5
+Pillow==8.4.0
+portalocker==2.5.1
+prometheus-client==0.12.0
+prompt-toolkit==3.0.22
+protobuf==3.19.1
+ptyprocess==0.7.0
+pyasn1==0.4.8
+pyasn1-modules==0.2.8
+pycocotools==2.0.6
+pycosat==0.6.3
+pycparser @ file:///tmp/build/80754af9/pycparser_1594388511720/work
+pydot==1.4.2
+Pygments==2.10.0
+pyOpenSSL @ file:///tmp/build/80754af9/pyopenssl_1608057966937/work
+pyparsing==3.0.6
+pyrsistent==0.18.0
+PySocks @ file:///tmp/build/80754af9/pysocks_1605305779399/work
+python-dateutil==2.8.2
+pytz==2021.3
+PyYAML==6.0
+pyzmq==22.3.0
+regex==2023.5.5
+requests @ file:///tmp/build/80754af9/requests_1608241421344/work
+requests-oauthlib==1.3.0
+rsa==4.8
+ruamel-yaml-conda @ file:///tmp/build/80754af9/ruamel_yaml_1616016699510/work
+Send2Trash==1.8.0
+setuptools-scm==6.3.2
+six @ file:///tmp/build/80754af9/six_1623709665295/work
+sniffio==1.2.0
+supervisor==4.2.2
+tabulate==0.8.10
+tensorboard==2.7.0
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.0
+termcolor==2.1.1
+terminado==0.12.1
+testpath==0.5.0
+toml==0.10.2
+tomli==1.2.2
+torch @ http://download.pytorch.org/whl/cu113/torch-1.10.0%2Bcu113-cp38-cp38-linux_x86_64.whl
+torchvision @ http://download.pytorch.org/whl/cu113/torchvision-0.11.1%2Bcu113-cp38-cp38-linux_x86_64.whl
+tornado==6.1
+tqdm @ file:///tmp/build/80754af9/tqdm_1625563689033/work
+traitlets==5.1.1
+typing-extensions==4.0.0
+urllib3 @ file:///tmp/build/80754af9/urllib3_1625084269274/work
+wcwidth==0.2.5
+webencodings==0.5.1
+websocket-client==1.2.1
+Werkzeug==2.0.2
+widgetsnbextension==3.5.2
+yacs==0.1.8
+zipp==3.6.0

detection/train.py

@@ -0,0 +1,251 @@
+data_root = '/root/autodl-tmp/ui_dataset'
+
+import logging
+import os
+from collections import OrderedDict
+import torch
+from torch.nn.parallel import DistributedDataParallel
+import random
+import cv2
+
+import detectron2.utils.comm as comm
+from detectron2.checkpoint import DetectionCheckpointer, PeriodicCheckpointer
+from detectron2.config import get_cfg
+from detectron2.utils.visualizer import Visualizer
+from detectron2.data import (
+    datasets,
+    MetadataCatalog,
+    get_detection_dataset_dicts,
+    build_detection_test_loader,
+    build_detection_train_loader,
+)
+from detectron2.engine import default_argument_parser, default_setup, default_writers, launch
+from detectron2.evaluation import (
+    CityscapesInstanceEvaluator,
+    CityscapesSemSegEvaluator,
+    COCOEvaluator,
+    COCOPanopticEvaluator,
+    DatasetEvaluators,
+    LVISEvaluator,
+    PascalVOCDetectionEvaluator,
+    SemSegEvaluator,
+    inference_on_dataset,
+    print_csv_format,
+)
+from detectron2.modeling import build_model
+from detectron2.solver import build_lr_scheduler, build_optimizer
+from detectron2.utils.events import EventStorage
+
+from icecream import ic, install
+install()
+ic.configureOutput(includeContext=True, contextAbsPath=True)
+
+logger = logging.getLogger("detectron2")
+
+
+def visualize(dataset_name='valid_ui', num=4, iter=0):
+    if not os.path.exists('./imgs'):
+        os.mkdir('./imgs')
+    metadata = MetadataCatalog.get(dataset_name)
+    dataset = get_detection_dataset_dicts(dataset_name)
+
+    for i, d in enumerate(random.sample(dataset, num)):
+        img = cv2.imread(d["file_name"])
+        visualizer = Visualizer(img[:, :, ::-1], metadata=metadata, scale=0.5)
+        vis = visualizer.draw_dataset_dict(d)
+        cv2.imwrite(f'./imgs/{iter}_{dataset_name}_{i}.png', vis.get_image()[:, :, ::-1])
+
+
+def get_evaluator(cfg, dataset_name, output_folder=None):
+    """
+    Create evaluator(s) for a given dataset.
+    This uses the special metadata "evaluator_type" associated with each builtin dataset.
+    For your own dataset, you can simply create an evaluator manually in your
+    script and do not have to worry about the hacky if-else logic here.
+    """
+    if output_folder is None:
+        output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
+    evaluator_list = []
+    evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
+    if evaluator_type in ["sem_seg", "coco_panoptic_seg"]:
+        evaluator_list.append(
+            SemSegEvaluator(
+                dataset_name,
+                distributed=True,
+                output_dir=output_folder,
+            )
+        )
+    if evaluator_type in ["coco", "coco_panoptic_seg"]:
+        evaluator_list.append(COCOEvaluator(dataset_name, output_dir=output_folder))
+    if evaluator_type == "coco_panoptic_seg":
+        evaluator_list.append(COCOPanopticEvaluator(dataset_name, output_folder))
+    if evaluator_type == "cityscapes_instance":
+        return CityscapesInstanceEvaluator(dataset_name)
+    if evaluator_type == "cityscapes_sem_seg":
+        return CityscapesSemSegEvaluator(dataset_name)
+    if evaluator_type == "pascal_voc":
+        return PascalVOCDetectionEvaluator(dataset_name)
+    if evaluator_type == "lvis":
+        return LVISEvaluator(dataset_name, cfg, True, output_folder)
+    if len(evaluator_list) == 0:
+        raise NotImplementedError(
+            "no Evaluator for the dataset {} with the type {}".format(dataset_name, evaluator_type)
+        )
+    if len(evaluator_list) == 1:
+        return evaluator_list[0]
+    return DatasetEvaluators(evaluator_list)
+
+
+def do_test(cfg, model, storage=None):
+    results = OrderedDict()
+    for dataset_name in cfg.DATASETS.TEST:
+        data_loader = build_detection_test_loader(cfg, dataset_name)
+        evaluator = get_evaluator(
+            cfg, dataset_name, os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
+        )
+        results_i = inference_on_dataset(model, data_loader, evaluator)
+        results[dataset_name] = results_i
+        if comm.is_main_process():
+            logger.info("Evaluation results for {} in csv format:".format(dataset_name))
+            print_csv_format(results_i)
+            # dump to storage, save to tensorboard
+            if storage != None:
+                for key, value in results_i.items(): # key = bbox / segm; value = {'AP': xx, 'APm': xx, ...}
+                    logging.info(f'key value: {key}, {value}')
+                    logging.info(f'key: {key}')
+                    out_aps_dict = {}
+                    for k, v in value.items():
+                        k = dataset_name + '_' + k
+                        out_aps_dict[k] = v
+                        # print('**{k: v.item() for k, v in comm.reduce_dict(results_i).items()}\n', type(**{k: v.item() for k, v in comm.reduce_dict(results_i).items()}))
+                    storage.put_scalars(**out_aps_dict)
+    if len(results) == 1:
+        results = list(results.values())[0]
+    return results
+
+
+def do_train(cfg, model, resume=False):
+    model.train()
+    optimizer = build_optimizer(cfg, model)
+    scheduler = build_lr_scheduler(cfg, optimizer)
+
+    checkpointer = DetectionCheckpointer(
+        model, cfg.OUTPUT_DIR, optimizer=optimizer, scheduler=scheduler
+    )
+    start_iter = (
+        checkpointer.resume_or_load(cfg.MODEL.WEIGHTS, resume=resume).get("iteration", -1) + 1
+    )
+    max_iter = cfg.SOLVER.MAX_ITER
+
+    periodic_checkpointer = PeriodicCheckpointer(
+        checkpointer, cfg.SOLVER.CHECKPOINT_PERIOD, max_iter=max_iter
+    )
+
+    writers = default_writers(cfg.OUTPUT_DIR, max_iter) if comm.is_main_process() else []
+
+    # compared to "train_net.py", we do not support accurate timing and
+    # precise BN here, because they are not trivial to implement in a small training loop
+    data_loader = build_detection_train_loader(cfg)
+    logger.info("Starting training from iteration {}".format(start_iter))
+    with EventStorage(start_iter) as storage:
+        for data, iteration in zip(data_loader, range(start_iter, max_iter)):
+            storage.iter = iteration
+
+            loss_dict = model(data)
+            losses = sum(loss_dict.values())
+            assert torch.isfinite(losses).all(), loss_dict
+
+            loss_dict_reduced = {k: v.item() for k, v in comm.reduce_dict(loss_dict).items()}
+            losses_reduced = sum(loss for loss in loss_dict_reduced.values())
+            if comm.is_main_process():
+                storage.put_scalars(total_loss=losses_reduced, **loss_dict_reduced)
+
+            optimizer.zero_grad()
+            losses.backward()
+            optimizer.step()
+            storage.put_scalar("lr", optimizer.param_groups[0]["lr"], smoothing_hint=False)
+            scheduler.step()
+
+            if (
+                cfg.TEST.EVAL_PERIOD > 0
+                and (iteration + 1) % cfg.TEST.EVAL_PERIOD == 0
+                and iteration != max_iter - 1
+            ):
+                visualize('valid_ui', 5, iteration)
+                visualize('train_ui', 5, iteration)
+                do_test(cfg, model, storage)
+                # Compared to "train_net.py", the test results are not dumped to EventStorage
+                comm.synchronize()
+
+            if iteration - start_iter > 5 and (
+                (iteration + 1) % 20 == 0 or iteration == max_iter - 1
+            ):
+                for writer in writers:
+                    writer.write()
+            periodic_checkpointer.step(iteration)
+
+
+def setup(args):
+    """
+    Create configs and perform basic setups.
+    """
+    cfg = get_cfg()
+    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+    default_setup(
+        cfg, args
+    )  # if you don't like any of the default setup, write your own setup code
+    return cfg
+
+
+def main(args):
+    cfg = setup(args)
+
+    datasets.register_coco_instances("train_ui", {},
+                                     f"{data_root}/train/_annotations.coco.json",
+                                     f"{data_root}/train")
+    datasets.register_coco_instances("train_dora_ui", {},
+                                     f"{data_root.replace('ui_dataset', 'dora_dataset')}/train.json",
+                                     f"{data_root.replace('ui_dataset', 'dora_dataset')}/train")
+    datasets.register_coco_instances("test_ui", {},
+                                     f"{data_root}/test/_annotations.coco.json",
+                                     f"{data_root}/test")
+    datasets.register_coco_instances("valid_ui", {},
+                                     f"{data_root}/valid/_annotations.coco.json",
+                                     f"{data_root}/valid")
+    datasets.register_coco_instances("valid_dora_ui", {},
+                                     f"{data_root.replace('ui_dataset', 'dora_dataset')}/val.json",
+                                     f"{data_root.replace('ui_dataset', 'dora_dataset')}/train")
+    print('done registering datasets')
+
+
+    model = build_model(cfg)
+    logger.info("Model:\n{}".format(model))
+    if args.eval_only:
+        DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
+            cfg.MODEL.WEIGHTS, resume=args.resume
+        )
+        return do_test(cfg, model)
+
+    distributed = comm.get_world_size() > 1
+    if distributed:
+        model = DistributedDataParallel(
+            model, device_ids=[comm.get_local_rank()], broadcast_buffers=False
+        )
+
+    do_train(cfg, model, resume=args.resume)
+    return do_test(cfg, model)
+
+
+if __name__ == "__main__":
+    args = default_argument_parser().parse_args()
+    print("Command Line Args:", args)
+    launch(
+        main,
+        args.num_gpus,
+        num_machines=args.num_machines,
+        machine_rank=args.machine_rank,
+        dist_url=args.dist_url,
+        args=(args,),
+    )