Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1# Copyright (c) Facebook, Inc. and its affiliates.2import logging3import os4from fvcore.common.timer import Timer5 6from detectron2.data import DatasetCatalog, MetadataCatalog7from detectron2.structures import BoxMode8from detectron2.utils.file_io import PathManager9 10from .builtin_meta import _get_coco_instances_meta11from .lvis_v0_5_categories import LVIS_CATEGORIES as LVIS_V0_5_CATEGORIES12from .lvis_v1_categories import LVIS_CATEGORIES as LVIS_V1_CATEGORIES13from .lvis_v1_category_image_count import LVIS_CATEGORY_IMAGE_COUNT as LVIS_V1_CATEGORY_IMAGE_COUNT14 15"""16This file contains functions to parse LVIS-format annotations into dicts in the17"Detectron2 format".18"""19 20logger = logging.getLogger(__name__)21 22__all__ = ["load_lvis_json", "register_lvis_instances", "get_lvis_instances_meta"]23 24 25def register_lvis_instances(name, metadata, json_file, image_root):26    """27    Register a dataset in LVIS's json annotation format for instance detection and segmentation.28 29    Args:30        name (str): a name that identifies the dataset, e.g. "lvis_v0.5_train".31        metadata (dict): extra metadata associated with this dataset. It can be an empty dict.32        json_file (str): path to the json instance annotation file.33        image_root (str or path-like): directory which contains all the images.34    """35    DatasetCatalog.register(name, lambda: load_lvis_json(json_file, image_root, name))36    MetadataCatalog.get(name).set(37        json_file=json_file, image_root=image_root, evaluator_type="lvis", **metadata38    )39 40 41def load_lvis_json(json_file, image_root, dataset_name=None, extra_annotation_keys=None):42    """43    Load a json file in LVIS's annotation format.44 45    Args:46        json_file (str): full path to the LVIS json annotation file.47        image_root (str): the directory where the images in this json file exists.48        dataset_name (str): the name of the dataset (e.g., "lvis_v0.5_train").49            If provided, this function will put "thing_classes" into the metadata50            associated with this dataset.51        extra_annotation_keys (list[str]): list of per-annotation keys that should also be52            loaded into the dataset dict (besides "bbox", "bbox_mode", "category_id",53            "segmentation"). The values for these keys will be returned as-is.54 55    Returns:56        list[dict]: a list of dicts in Detectron2 standard format. (See57        `Using Custom Datasets </tutorials/datasets.html>`_ )58 59    Notes:60        1. This function does not read the image files.61           The results do not have the "image" field.62    """63    from lvis import LVIS64 65    json_file = PathManager.get_local_path(json_file)66 67    timer = Timer()68    lvis_api = LVIS(json_file)69    if timer.seconds() > 1:70        logger.info("Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds()))71 72    if dataset_name is not None:73        meta = get_lvis_instances_meta(dataset_name)74        MetadataCatalog.get(dataset_name).set(**meta)75 76    # sort indices for reproducible results77    img_ids = sorted(lvis_api.imgs.keys())78    # imgs is a list of dicts, each looks something like:79    # {'license': 4,80    #  'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg',81    #  'file_name': 'COCO_val2014_000000001268.jpg',82    #  'height': 427,83    #  'width': 640,84    #  'date_captured': '2013-11-17 05:57:24',85    #  'id': 1268}86    imgs = lvis_api.load_imgs(img_ids)87    # anns is a list[list[dict]], where each dict is an annotation88    # record for an object. The inner list enumerates the objects in an image89    # and the outer list enumerates over images. Example of anns[0]:90    # [{'segmentation': [[192.81,91    #     247.09,92    #     ...93    #     219.03,94    #     249.06]],95    #   'area': 1035.749,96    #   'image_id': 1268,97    #   'bbox': [192.81, 224.8, 74.73, 33.43],98    #   'category_id': 16,99    #   'id': 42986},100    #  ...]101    anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]102 103    # Sanity check that each annotation has a unique id104    ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]105    assert len(set(ann_ids)) == len(ann_ids), "Annotation ids in '{}' are not unique".format(106        json_file107    )108 109    imgs_anns = list(zip(imgs, anns))110 111    logger.info("Loaded {} images in the LVIS format from {}".format(len(imgs_anns), json_file))112 113    if extra_annotation_keys:114        logger.info(115            "The following extra annotation keys will be loaded: {} ".format(extra_annotation_keys)116        )117    else:118        extra_annotation_keys = []119 120    def get_file_name(img_root, img_dict):121        # Determine the path including the split folder ("train2017", "val2017", "test2017") from122        # the coco_url field. Example:123        #   'coco_url': 'http://images.cocodataset.org/train2017/000000155379.jpg'124        split_folder, file_name = img_dict["coco_url"].split("/")[-2:]125        return os.path.join(img_root + split_folder, file_name)126 127    dataset_dicts = []128 129    for (img_dict, anno_dict_list) in imgs_anns:130        record = {}131        record["file_name"] = get_file_name(image_root, img_dict)132        record["height"] = img_dict["height"]133        record["width"] = img_dict["width"]134        record["not_exhaustive_category_ids"] = img_dict.get("not_exhaustive_category_ids", [])135        record["neg_category_ids"] = img_dict.get("neg_category_ids", [])136        image_id = record["image_id"] = img_dict["id"]137 138        objs = []139        for anno in anno_dict_list:140            # Check that the image_id in this annotation is the same as141            # the image_id we're looking at.142            # This fails only when the data parsing logic or the annotation file is buggy.143            assert anno["image_id"] == image_id144            obj = {"bbox": anno["bbox"], "bbox_mode": BoxMode.XYWH_ABS}145            # LVIS data loader can be used to load COCO dataset categories. In this case `meta`146            # variable will have a field with COCO-specific category mapping.147            if dataset_name is not None and "thing_dataset_id_to_contiguous_id" in meta:148                obj["category_id"] = meta["thing_dataset_id_to_contiguous_id"][anno["category_id"]]149            else:150                obj["category_id"] = anno["category_id"] - 1  # Convert 1-indexed to 0-indexed151            segm = anno["segmentation"]  # list[list[float]]152            # filter out invalid polygons (< 3 points)153            valid_segm = [poly for poly in segm if len(poly) % 2 == 0 and len(poly) >= 6]154            assert len(segm) == len(155                valid_segm156            ), "Annotation contains an invalid polygon with < 3 points"157            assert len(segm) > 0158            obj["segmentation"] = segm159            for extra_ann_key in extra_annotation_keys:160                obj[extra_ann_key] = anno[extra_ann_key]161            objs.append(obj)162        record["annotations"] = objs163        dataset_dicts.append(record)164 165    return dataset_dicts166 167 168def get_lvis_instances_meta(dataset_name):169    """170    Load LVIS metadata.171 172    Args:173        dataset_name (str): LVIS dataset name without the split name (e.g., "lvis_v0.5").174 175    Returns:176        dict: LVIS metadata with keys: thing_classes177    """178    if "cocofied" in dataset_name:179        return _get_coco_instances_meta()180    if "v0.5" in dataset_name:181        return _get_lvis_instances_meta_v0_5()182    elif "v1" in dataset_name:183        return _get_lvis_instances_meta_v1()184    raise ValueError("No built-in metadata for dataset {}".format(dataset_name))185 186 187def _get_lvis_instances_meta_v0_5():188    assert len(LVIS_V0_5_CATEGORIES) == 1230189    cat_ids = [k["id"] for k in LVIS_V0_5_CATEGORIES]190    assert min(cat_ids) == 1 and max(cat_ids) == len(191        cat_ids192    ), "Category ids are not in [1, #categories], as expected"193    # Ensure that the category list is sorted by id194    lvis_categories = sorted(LVIS_V0_5_CATEGORIES, key=lambda x: x["id"])195    thing_classes = [k["synonyms"][0] for k in lvis_categories]196    meta = {"thing_classes": thing_classes}197    return meta198 199 200def _get_lvis_instances_meta_v1():201    assert len(LVIS_V1_CATEGORIES) == 1203202    cat_ids = [k["id"] for k in LVIS_V1_CATEGORIES]203    assert min(cat_ids) == 1 and max(cat_ids) == len(204        cat_ids205    ), "Category ids are not in [1, #categories], as expected"206    # Ensure that the category list is sorted by id207    lvis_categories = sorted(LVIS_V1_CATEGORIES, key=lambda x: x["id"])208    thing_classes = [k["synonyms"][0] for k in lvis_categories]209    meta = {"thing_classes": thing_classes, "class_image_count": LVIS_V1_CATEGORY_IMAGE_COUNT}210    return meta211 212 213if __name__ == "__main__":214    """215    Test the LVIS json dataset loader.216 217    Usage:218        python -m detectron2.data.datasets.lvis \219            path/to/json path/to/image_root dataset_name vis_limit220    """221    import sys222    import numpy as np223    from detectron2.utils.logger import setup_logger224    from PIL import Image225    import detectron2.data.datasets  # noqa # add pre-defined metadata226    from detectron2.utils.visualizer import Visualizer227 228    logger = setup_logger(name=__name__)229    meta = MetadataCatalog.get(sys.argv[3])230 231    dicts = load_lvis_json(sys.argv[1], sys.argv[2], sys.argv[3])232    logger.info("Done loading {} samples.".format(len(dicts)))233 234    dirname = "lvis-data-vis"235    os.makedirs(dirname, exist_ok=True)236    for d in dicts[: int(sys.argv[4])]:237        img = np.array(Image.open(d["file_name"]))238        visualizer = Visualizer(img, metadata=meta)239        vis = visualizer.draw_dataset_dict(d)240        fpath = os.path.join(dirname, os.path.basename(d["file_name"]))241        vis.save(fpath)242