Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1# Copyright (c) Facebook, Inc. and its affiliates.2import itertools3import json4import logging5import numpy as np6import os7from collections import OrderedDict8from typing import Optional, Union9import pycocotools.mask as mask_util10import torch11from PIL import Image12 13from detectron2.data import DatasetCatalog, MetadataCatalog14from detectron2.utils.comm import all_gather, is_main_process, synchronize15from detectron2.utils.file_io import PathManager16 17from .evaluator import DatasetEvaluator18 19_CV2_IMPORTED = True20try:21    import cv2  # noqa22except ImportError:23    # OpenCV is an optional dependency at the moment24    _CV2_IMPORTED = False25 26 27def load_image_into_numpy_array(28    filename: str,29    copy: bool = False,30    dtype: Optional[Union[np.dtype, str]] = None,31) -> np.ndarray:32    with PathManager.open(filename, "rb") as f:33        array = np.array(Image.open(f), copy=copy, dtype=dtype)34    return array35 36 37class SemSegEvaluator(DatasetEvaluator):38    """39    Evaluate semantic segmentation metrics.40    """41 42    def __init__(43        self,44        dataset_name,45        distributed=True,46        output_dir=None,47        *,48        sem_seg_loading_fn=load_image_into_numpy_array,49        num_classes=None,50        ignore_label=None,51    ):52        """53        Args:54            dataset_name (str): name of the dataset to be evaluated.55            distributed (bool): if True, will collect results from all ranks for evaluation.56                Otherwise, will evaluate the results in the current process.57            output_dir (str): an output directory to dump results.58            sem_seg_loading_fn: function to read sem seg file and load into numpy array.59                Default provided, but projects can customize.60            num_classes, ignore_label: deprecated argument61        """62        self._logger = logging.getLogger(__name__)63        if num_classes is not None:64            self._logger.warn(65                "SemSegEvaluator(num_classes) is deprecated! It should be obtained from metadata."66            )67        if ignore_label is not None:68            self._logger.warn(69                "SemSegEvaluator(ignore_label) is deprecated! It should be obtained from metadata."70            )71        self._dataset_name = dataset_name72        self._distributed = distributed73        self._output_dir = output_dir74 75        self._cpu_device = torch.device("cpu")76 77        self.input_file_to_gt_file = {78            dataset_record["file_name"]: dataset_record["sem_seg_file_name"]79            for dataset_record in DatasetCatalog.get(dataset_name)80        }81 82        meta = MetadataCatalog.get(dataset_name)83        # Dict that maps contiguous training ids to COCO category ids84        try:85            c2d = meta.stuff_dataset_id_to_contiguous_id86            self._contiguous_id_to_dataset_id = {v: k for k, v in c2d.items()}87        except AttributeError:88            self._contiguous_id_to_dataset_id = None89        self._class_names = meta.stuff_classes90        self.sem_seg_loading_fn = sem_seg_loading_fn91        self._num_classes = len(meta.stuff_classes)92        if num_classes is not None:93            assert self._num_classes == num_classes, f"{self._num_classes} != {num_classes}"94        self._ignore_label = ignore_label if ignore_label is not None else meta.ignore_label95 96        # This is because cv2.erode did not work for int datatype. Only works for uint8.97        self._compute_boundary_iou = True98        if not _CV2_IMPORTED:99            self._compute_boundary_iou = False100            self._logger.warn(101                """Boundary IoU calculation requires OpenCV. B-IoU metrics are102                not going to be computed because OpenCV is not available to import."""103            )104        if self._num_classes >= np.iinfo(np.uint8).max:105            self._compute_boundary_iou = False106            self._logger.warn(107                f"""SemSegEvaluator(num_classes) is more than supported value for Boundary IoU calculation!108                B-IoU metrics are not going to be computed. Max allowed value (exclusive)109                for num_classes for calculating Boundary IoU is {np.iinfo(np.uint8).max}.110                The number of classes of dataset {self._dataset_name} is {self._num_classes}"""111            )112 113    def reset(self):114        self._conf_matrix = np.zeros((self._num_classes + 1, self._num_classes + 1), dtype=np.int64)115        self._b_conf_matrix = np.zeros(116            (self._num_classes + 1, self._num_classes + 1), dtype=np.int64117        )118        self._predictions = []119 120    def process(self, inputs, outputs):121        """122        Args:123            inputs: the inputs to a model.124                It is a list of dicts. Each dict corresponds to an image and125                contains keys like "height", "width", "file_name".126            outputs: the outputs of a model. It is either list of semantic segmentation predictions127                (Tensor [H, W]) or list of dicts with key "sem_seg" that contains semantic128                segmentation prediction in the same format.129        """130        for input, output in zip(inputs, outputs):131            output = output["sem_seg"].argmax(dim=0).to(self._cpu_device)132            pred = np.array(output, dtype=np.int)133            gt_filename = self.input_file_to_gt_file[input["file_name"]]134            gt = self.sem_seg_loading_fn(gt_filename, dtype=np.int)135 136            gt[gt == self._ignore_label] = self._num_classes137 138            self._conf_matrix += np.bincount(139                (self._num_classes + 1) * pred.reshape(-1) + gt.reshape(-1),140                minlength=self._conf_matrix.size,141            ).reshape(self._conf_matrix.shape)142 143            if self._compute_boundary_iou:144                b_gt = self._mask_to_boundary(gt.astype(np.uint8))145                b_pred = self._mask_to_boundary(pred.astype(np.uint8))146 147                self._b_conf_matrix += np.bincount(148                    (self._num_classes + 1) * b_pred.reshape(-1) + b_gt.reshape(-1),149                    minlength=self._conf_matrix.size,150                ).reshape(self._conf_matrix.shape)151 152            self._predictions.extend(self.encode_json_sem_seg(pred, input["file_name"]))153 154    def evaluate(self):155        """156        Evaluates standard semantic segmentation metrics (http://cocodataset.org/#stuff-eval):157 158        * Mean intersection-over-union averaged across classes (mIoU)159        * Frequency Weighted IoU (fwIoU)160        * Mean pixel accuracy averaged across classes (mACC)161        * Pixel Accuracy (pACC)162        """163        if self._distributed:164            synchronize()165            conf_matrix_list = all_gather(self._conf_matrix)166            b_conf_matrix_list = all_gather(self._b_conf_matrix)167            self._predictions = all_gather(self._predictions)168            self._predictions = list(itertools.chain(*self._predictions))169            if not is_main_process():170                return171 172            self._conf_matrix = np.zeros_like(self._conf_matrix)173            for conf_matrix in conf_matrix_list:174                self._conf_matrix += conf_matrix175 176            self._b_conf_matrix = np.zeros_like(self._b_conf_matrix)177            for b_conf_matrix in b_conf_matrix_list:178                self._b_conf_matrix += b_conf_matrix179 180        if self._output_dir:181            PathManager.mkdirs(self._output_dir)182            file_path = os.path.join(self._output_dir, "sem_seg_predictions.json")183            with PathManager.open(file_path, "w") as f:184                f.write(json.dumps(self._predictions))185 186        acc = np.full(self._num_classes, np.nan, dtype=np.float)187        iou = np.full(self._num_classes, np.nan, dtype=np.float)188        tp = self._conf_matrix.diagonal()[:-1].astype(np.float)189        pos_gt = np.sum(self._conf_matrix[:-1, :-1], axis=0).astype(np.float)190        class_weights = pos_gt / np.sum(pos_gt)191        pos_pred = np.sum(self._conf_matrix[:-1, :-1], axis=1).astype(np.float)192        acc_valid = pos_gt > 0193        acc[acc_valid] = tp[acc_valid] / pos_gt[acc_valid]194        union = pos_gt + pos_pred - tp195        iou_valid = np.logical_and(acc_valid, union > 0)196        iou[iou_valid] = tp[iou_valid] / union[iou_valid]197        macc = np.sum(acc[acc_valid]) / np.sum(acc_valid)198        miou = np.sum(iou[iou_valid]) / np.sum(iou_valid)199        fiou = np.sum(iou[iou_valid] * class_weights[iou_valid])200        pacc = np.sum(tp) / np.sum(pos_gt)201 202        if self._compute_boundary_iou:203            b_iou = np.full(self._num_classes, np.nan, dtype=np.float)204            b_tp = self._b_conf_matrix.diagonal()[:-1].astype(np.float)205            b_pos_gt = np.sum(self._b_conf_matrix[:-1, :-1], axis=0).astype(np.float)206            b_pos_pred = np.sum(self._b_conf_matrix[:-1, :-1], axis=1).astype(np.float)207            b_union = b_pos_gt + b_pos_pred - b_tp208            b_iou_valid = b_union > 0209            b_iou[b_iou_valid] = b_tp[b_iou_valid] / b_union[b_iou_valid]210 211        res = {}212        res["mIoU"] = 100 * miou213        res["fwIoU"] = 100 * fiou214        for i, name in enumerate(self._class_names):215            res[f"IoU-{name}"] = 100 * iou[i]216            if self._compute_boundary_iou:217                res[f"BoundaryIoU-{name}"] = 100 * b_iou[i]218                res[f"min(IoU, B-Iou)-{name}"] = 100 * min(iou[i], b_iou[i])219        res["mACC"] = 100 * macc220        res["pACC"] = 100 * pacc221        for i, name in enumerate(self._class_names):222            res[f"ACC-{name}"] = 100 * acc[i]223 224        if self._output_dir:225            file_path = os.path.join(self._output_dir, "sem_seg_evaluation.pth")226            with PathManager.open(file_path, "wb") as f:227                torch.save(res, f)228        results = OrderedDict({"sem_seg": res})229        self._logger.info(results)230        return results231 232    def encode_json_sem_seg(self, sem_seg, input_file_name):233        """234        Convert semantic segmentation to COCO stuff format with segments encoded as RLEs.235        See http://cocodataset.org/#format-results236        """237        json_list = []238        for label in np.unique(sem_seg):239            if self._contiguous_id_to_dataset_id is not None:240                assert (241                    label in self._contiguous_id_to_dataset_id242                ), "Label {} is not in the metadata info for {}".format(label, self._dataset_name)243                dataset_id = self._contiguous_id_to_dataset_id[label]244            else:245                dataset_id = int(label)246            mask = (sem_seg == label).astype(np.uint8)247            mask_rle = mask_util.encode(np.array(mask[:, :, None], order="F"))[0]248            mask_rle["counts"] = mask_rle["counts"].decode("utf-8")249            json_list.append(250                {"file_name": input_file_name, "category_id": dataset_id, "segmentation": mask_rle}251            )252        return json_list253 254    def _mask_to_boundary(self, mask: np.ndarray, dilation_ratio=0.02):255        assert mask.ndim == 2, "mask_to_boundary expects a 2-dimensional image"256        h, w = mask.shape257        diag_len = np.sqrt(h**2 + w**2)258        dilation = max(1, int(round(dilation_ratio * diag_len)))259        kernel = np.ones((3, 3), dtype=np.uint8)260 261        padded_mask = cv2.copyMakeBorder(mask, 1, 1, 1, 1, cv2.BORDER_CONSTANT, value=0)262        eroded_mask_with_padding = cv2.erode(padded_mask, kernel, iterations=dilation)263        eroded_mask = eroded_mask_with_padding[1:-1, 1:-1]264        boundary = mask - eroded_mask265        return boundary266