Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1# -*- coding: utf-8 -*-2# Copyright (c) Facebook, Inc. and its affiliates.3 4import numpy as np5import os6import xml.etree.ElementTree as ET7from typing import List, Tuple, Union8 9from detectron2.data import DatasetCatalog, MetadataCatalog10from detectron2.structures import BoxMode11from detectron2.utils.file_io import PathManager12 13__all__ = ["load_voc_instances", "register_pascal_voc"]14 15 16# fmt: off17CLASS_NAMES = (18    "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat",19    "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person",20    "pottedplant", "sheep", "sofa", "train", "tvmonitor"21)22# fmt: on23 24 25def load_voc_instances(dirname: str, split: str, class_names: Union[List[str], Tuple[str, ...]]):26    """27    Load Pascal VOC detection annotations to Detectron2 format.28 29    Args:30        dirname: Contain "Annotations", "ImageSets", "JPEGImages"31        split (str): one of "train", "test", "val", "trainval"32        class_names: list or tuple of class names33    """34    with PathManager.open(os.path.join(dirname, "ImageSets", "Main", split + ".txt")) as f:35        fileids = np.loadtxt(f, dtype=np.str)36 37    # Needs to read many small annotation files. Makes sense at local38    annotation_dirname = PathManager.get_local_path(os.path.join(dirname, "Annotations/"))39    dicts = []40    for fileid in fileids:41        anno_file = os.path.join(annotation_dirname, fileid + ".xml")42        jpeg_file = os.path.join(dirname, "JPEGImages", fileid + ".jpg")43 44        with PathManager.open(anno_file) as f:45            tree = ET.parse(f)46 47        r = {48            "file_name": jpeg_file,49            "image_id": fileid,50            "height": int(tree.findall("./size/height")[0].text),51            "width": int(tree.findall("./size/width")[0].text),52        }53        instances = []54 55        for obj in tree.findall("object"):56            cls = obj.find("name").text57            # We include "difficult" samples in training.58            # Based on limited experiments, they don't hurt accuracy.59            # difficult = int(obj.find("difficult").text)60            # if difficult == 1:61            # continue62            bbox = obj.find("bndbox")63            bbox = [float(bbox.find(x).text) for x in ["xmin", "ymin", "xmax", "ymax"]]64            # Original annotations are integers in the range [1, W or H]65            # Assuming they mean 1-based pixel indices (inclusive),66            # a box with annotation (xmin=1, xmax=W) covers the whole image.67            # In coordinate space this is represented by (xmin=0, xmax=W)68            bbox[0] -= 1.069            bbox[1] -= 1.070            instances.append(71                {"category_id": class_names.index(cls), "bbox": bbox, "bbox_mode": BoxMode.XYXY_ABS}72            )73        r["annotations"] = instances74        dicts.append(r)75    return dicts76 77 78def register_pascal_voc(name, dirname, split, year, class_names=CLASS_NAMES):79    DatasetCatalog.register(name, lambda: load_voc_instances(dirname, split, class_names))80    MetadataCatalog.get(name).set(81        thing_classes=list(class_names), dirname=dirname, year=year, split=split82    )83