from pycocotools import mask as maskUtils
import mmcv
import numpy as np

from .coco import CocoDataset
from .builder import DATASETS

@DATASETS.register_module()
class OCHumanDataset(CocoDataset):
  CLASSES = ('person')

  def _ochuman_segm2json(self, results):
    """Convert instance segmentation results to COCO json style."""
    bbox_json_results = []
    segm_json_results = []
    for idx in range(len(self)):
        img_id = self.img_ids[idx]
        det, seg = results[idx]
        for label in range(len(det)):
            # bbox results
            bboxes = det[label]
            for i in range(bboxes.shape[0]):
                data = dict()
                data['image_id'] = img_id
                data['bbox'] = self.xyxy2xywh(bboxes[i])
                data['score'] = float(bboxes[i][4])
                data['category_id'] = self.cat_ids[label]
                bbox_json_results.append(data)

            # segm results
            # some detectors use different scores for bbox and mask
            if isinstance(seg, tuple):
                segms = seg[0][label]
                mask_score = seg[1][label]
            else:
                segms = seg[label]
                mask_score = [bbox[4] for bbox in bboxes]
            for i in range(bboxes.shape[0]):
                data = dict()
                data['image_id'] = img_id
                data['bbox'] = self.xyxy2xywh(bboxes[i])
                data['score'] = float(mask_score[i])
                data['category_id'] = self.cat_ids[label]
                maskencode = maskUtils.encode(np.asfortranarray(segms[i]))
                maskencode['counts'] = maskencode['counts'].decode('ascii')
                data['segmentation'] = segms[i]
                segm_json_results.append(data)
        return bbox_json_results, segm_json_results

    def results2json(self, results, outfile_prefix):
        """Dump the detection results to a COCO style json file.
        There are 3 types of results: proposals, bbox predictions, mask
        predictions, and they have different data types. This method will
        automatically recognize the type, and dump them to json files.
        Args:
            results (list[list | tuple | ndarray]): Testing results of the
                dataset.
            outfile_prefix (str): The filename prefix of the json files. If the
                prefix is "somepath/xxx", the json files will be named
                "somepath/xxx.bbox.json", "somepath/xxx.segm.json",
                "somepath/xxx.proposal.json".
        Returns:
            dict[str: str]: Possible keys are "bbox", "segm", "proposal", and \
                values are corresponding filenames.
        """
        result_files = dict()
        if isinstance(results[0], list):
            json_results = self._det2json(results)
            result_files['bbox'] = f'{outfile_prefix}.bbox.json'
            result_files['proposal'] = f'{outfile_prefix}.bbox.json'
            mmcv.dump(json_results, result_files['bbox'])
        elif isinstance(results[0], tuple):
            json_results = self._ochuman_segm2json(results)
            result_files['bbox'] = f'{outfile_prefix}.bbox.json'
            result_files['proposal'] = f'{outfile_prefix}.bbox.json'
            result_files['segm'] = f'{outfile_prefix}.segm.json'
            mmcv.dump(json_results[0], result_files['bbox'])
            mmcv.dump(json_results[1], result_files['segm'])
        elif isinstance(results[0], np.ndarray):
            json_results = self._proposal2json(results)
            result_files['proposal'] = f'{outfile_prefix}.proposal.json'
            mmcv.dump(json_results, result_files['proposal'])
        else:
            raise TypeError('invalid type of results')
        return result_files

from pycocotools import mask as maskUtils
import mmcv
import numpy as np

from .coco import CocoDataset
from .builder import DATASETS

@DATASETS.register_module()
class OCHumanDataset(CocoDataset):
  CLASSES = ('person')

  def _ochuman_segm2json(self, results):
    """Convert instance segmentation results to COCO json style."""
    bbox_json_results = []
    segm_json_results = []
    for idx in range(len(self)):
        img_id = self.img_ids[idx]
        det, seg = results[idx]
        for label in range(len(det)):
            # bbox results
            bboxes = det[label]
            for i in range(bboxes.shape[0]):
                data = dict()
                data['image_id'] = img_id
                data['bbox'] = self.xyxy2xywh(bboxes[i])
                data['score'] = float(bboxes[i][4])
                data['category_id'] = self.cat_ids[label]
                bbox_json_results.append(data)

            # segm results
            # some detectors use different scores for bbox and mask
            if isinstance(seg, tuple):
                segms = seg[0][label]
                mask_score = seg[1][label]
            else:
                segms = seg[label]
                mask_score = [bbox[4] for bbox in bboxes]
            for i in range(bboxes.shape[0]):
                data = dict()
                data['image_id'] = img_id
                data['bbox'] = self.xyxy2xywh(bboxes[i])
                data['score'] = float(mask_score[i])
                data['category_id'] = self.cat_ids[label]
                maskencode = maskUtils.encode(np.asfortranarray(segms[i]))
                maskencode['counts'] = maskencode['counts'].decode('ascii')
                data['segmentation'] = segms[i]
                segm_json_results.append(data)
        return bbox_json_results, segm_json_results

    def results2json(self, results, outfile_prefix):
        """Dump the detection results to a COCO style json file.
        There are 3 types of results: proposals, bbox predictions, mask
        predictions, and they have different data types. This method will
        automatically recognize the type, and dump them to json files.
        Args:
            results (list[list | tuple | ndarray]): Testing results of the
                dataset.
            outfile_prefix (str): The filename prefix of the json files. If the
                prefix is "somepath/xxx", the json files will be named
                "somepath/xxx.bbox.json", "somepath/xxx.segm.json",
                "somepath/xxx.proposal.json".
        Returns:
            dict[str: str]: Possible keys are "bbox", "segm", "proposal", and \
                values are corresponding filenames.
        """
        result_files = dict()
        if isinstance(results[0], list):
            json_results = self._det2json(results)
            result_files['bbox'] = f'{outfile_prefix}.bbox.json'
            result_files['proposal'] = f'{outfile_prefix}.bbox.json'
            mmcv.dump(json_results, result_files['bbox'])
        elif isinstance(results[0], tuple):
            json_results = self._ochuman_segm2json(results)
            result_files['bbox'] = f'{outfile_prefix}.bbox.json'
            result_files['proposal'] = f'{outfile_prefix}.bbox.json'
            result_files['segm'] = f'{outfile_prefix}.segm.json'
            mmcv.dump(json_results[0], result_files['bbox'])
            mmcv.dump(json_results[1], result_files['segm'])
        elif isinstance(results[0], np.ndarray):
            json_results = self._proposal2json(results)
            result_files['proposal'] = f'{outfile_prefix}.proposal.json'
            mmcv.dump(json_results, result_files['proposal'])
        else:
            raise TypeError('invalid type of results')
        return result_files

Convert instance segmentation results to COCO json style. # bbox results # segm results # some detectors use different scores for bbox and mask Dump the detection results to a COCO style json file.         There are 3 types of results: proposals, bbox predictions, mask         predictions, and they have different data types. This method will         automatically recognize the type, and dump them to json files.         Args:             results (list[list | tuple | ndarray]): Testing results of the                 dataset.             outfile_prefix (str): The filename prefix of the json files. If the                 prefix is "somepath/xxx", the json files will be named                 "somepath/xxx.bbox.json", "somepath/xxx.segm.json",                 "somepath/xxx.proposal.json".         Returns:             dict[str: str]: Possible keys are "bbox", "segm", "proposal", and \                 values are corresponding filenames.