Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1#!/usr/bin/env python32# -*- coding: utf-8 -*-3# Copyright (c) Facebook, Inc. and its affiliates.4 5import functools6import json7import multiprocessing as mp8import numpy as np9import os10import time11from fvcore.common.download import download12from panopticapi.utils import rgb2id13from PIL import Image14 15from detectron2.data.datasets.builtin_meta import COCO_CATEGORIES16 17 18def _process_panoptic_to_semantic(input_panoptic, output_semantic, segments, id_map):19    panoptic = np.asarray(Image.open(input_panoptic), dtype=np.uint32)20    panoptic = rgb2id(panoptic)21    output = np.zeros_like(panoptic, dtype=np.uint8) + 25522    for seg in segments:23        cat_id = seg["category_id"]24        new_cat_id = id_map[cat_id]25        output[panoptic == seg["id"]] = new_cat_id26    Image.fromarray(output).save(output_semantic)27 28 29def separate_coco_semantic_from_panoptic(panoptic_json, panoptic_root, sem_seg_root, categories):30    """31    Create semantic segmentation annotations from panoptic segmentation32    annotations, to be used by PanopticFPN.33 34    It maps all thing categories to class 0, and maps all unlabeled pixels to class 255.35    It maps all stuff categories to contiguous ids starting from 1.36 37    Args:38        panoptic_json (str): path to the panoptic json file, in COCO's format.39        panoptic_root (str): a directory with panoptic annotation files, in COCO's format.40        sem_seg_root (str): a directory to output semantic annotation files41        categories (list[dict]): category metadata. Each dict needs to have:42            "id": corresponds to the "category_id" in the json annotations43            "isthing": 0 or 144    """45    os.makedirs(sem_seg_root, exist_ok=True)46 47    stuff_ids = [k["id"] for k in categories if k["isthing"] == 0]48    thing_ids = [k["id"] for k in categories if k["isthing"] == 1]49    id_map = {}  # map from category id to id in the output semantic annotation50    assert len(stuff_ids) <= 25451    for i, stuff_id in enumerate(stuff_ids):52        id_map[stuff_id] = i + 153    for thing_id in thing_ids:54        id_map[thing_id] = 055    id_map[0] = 25556 57    with open(panoptic_json) as f:58        obj = json.load(f)59 60    pool = mp.Pool(processes=max(mp.cpu_count() // 2, 4))61 62    def iter_annotations():63        for anno in obj["annotations"]:64            file_name = anno["file_name"]65            segments = anno["segments_info"]66            input = os.path.join(panoptic_root, file_name)67            output = os.path.join(sem_seg_root, file_name)68            yield input, output, segments69 70    print("Start writing to {} ...".format(sem_seg_root))71    start = time.time()72    pool.starmap(73        functools.partial(_process_panoptic_to_semantic, id_map=id_map),74        iter_annotations(),75        chunksize=100,76    )77    print("Finished. time: {:.2f}s".format(time.time() - start))78 79 80if __name__ == "__main__":81    dataset_dir = os.path.join(os.getenv("DETECTRON2_DATASETS", "datasets"), "coco")82    for s in ["val2017", "train2017"]:83        separate_coco_semantic_from_panoptic(84            os.path.join(dataset_dir, "annotations/panoptic_{}.json".format(s)),85            os.path.join(dataset_dir, "panoptic_{}".format(s)),86            os.path.join(dataset_dir, "panoptic_stuff_{}".format(s)),87            COCO_CATEGORIES,88        )89 90    # Prepare val2017_100 for quick testing:91 92    dest_dir = os.path.join(dataset_dir, "annotations/")93    URL_PREFIX = "https://dl.fbaipublicfiles.com/detectron2/"94    download(URL_PREFIX + "annotations/coco/panoptic_val2017_100.json", dest_dir)95    with open(os.path.join(dest_dir, "panoptic_val2017_100.json")) as f:96        obj = json.load(f)97 98    def link_val100(dir_full, dir_100):99        print("Creating " + dir_100 + " ...")100        os.makedirs(dir_100, exist_ok=True)101        for img in obj["images"]:102            basename = os.path.splitext(img["file_name"])[0]103            src = os.path.join(dir_full, basename + ".png")104            dst = os.path.join(dir_100, basename + ".png")105            src = os.path.relpath(src, start=dir_100)106            os.symlink(src, dst)107 108    link_val100(109        os.path.join(dataset_dir, "panoptic_val2017"),110        os.path.join(dataset_dir, "panoptic_val2017_100"),111    )112 113    link_val100(114        os.path.join(dataset_dir, "panoptic_stuff_val2017"),115        os.path.join(dataset_dir, "panoptic_stuff_val2017_100"),116    )117