maskrcnn_benchmark/data/datasets/vg.py

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import collections
import json
import os.path as op

import numpy as np
import torch

from .tsv import TSVYamlDataset, find_file_path_in_yaml
from .box_label_loader import BoxLabelLoader
from maskrcnn_benchmark.data.datasets.coco_dt import CocoDetectionTSV


class VGDetectionTSV(CocoDetectionTSV):
    pass


def sort_key_by_val(dic):
    sorted_dic = sorted(dic.items(), key=lambda kv: kv[1])
    return [kv[0] for kv in sorted_dic]


def bbox_overlaps(anchors, gt_boxes):
    """
    anchors: (N, 4) ndarray of float
    gt_boxes: (K, 4) ndarray of float
    overlaps: (N, K) ndarray of overlap between boxes and query_boxes
    """
    N = anchors.size(0)
    K = gt_boxes.size(0)

    gt_boxes_area = ((gt_boxes[:, 2] - gt_boxes[:, 0] + 1) *
                     (gt_boxes[:, 3] - gt_boxes[:, 1] + 1)).view(1, K)

    anchors_area = ((anchors[:, 2] - anchors[:, 0] + 1) *
                    (anchors[:, 3] - anchors[:, 1] + 1)).view(N, 1)

    boxes = anchors.view(N, 1, 4).expand(N, K, 4)
    query_boxes = gt_boxes.view(1, K, 4).expand(N, K, 4)

    iw = (torch.min(boxes[:, :, 2], query_boxes[:, :, 2]) -
          torch.max(boxes[:, :, 0], query_boxes[:, :, 0]) + 1)
    iw[iw < 0] = 0

    ih = (torch.min(boxes[:, :, 3], query_boxes[:, :, 3]) -
          torch.max(boxes[:, :, 1], query_boxes[:, :, 1]) + 1)
    ih[ih < 0] = 0

    ua = anchors_area + gt_boxes_area - (iw * ih)
    overlaps = iw * ih / ua

    return overlaps


# VG data loader for Danfei Xu's Scene graph focused format.
# todo: if ordering of classes, attributes, relations changed
# todo make sure to re-write the obj_classes.txt/rel_classes.txt files

def _box_filter(boxes, must_overlap=False):
    """ Only include boxes that overlap as possible relations.
    If no overlapping boxes, use all of them."""
    overlaps = bbox_overlaps(boxes, boxes).numpy() > 0
    np.fill_diagonal(overlaps, 0)

    all_possib = np.ones_like(overlaps, dtype=np.bool)
    np.fill_diagonal(all_possib, 0)

    if must_overlap:
        possible_boxes = np.column_stack(np.where(overlaps))

        if possible_boxes.size == 0:
            possible_boxes = np.column_stack(np.where(all_possib))
    else:
        possible_boxes = np.column_stack(np.where(all_possib))
    return possible_boxes


class VGTSVDataset(TSVYamlDataset):
    """
    Generic TSV dataset format for Object Detection.
    """

    def __init__(self, yaml_file, extra_fields=None, transforms=None,
                 is_load_label=True, filter_duplicate_rels=True,
                 relation_on=False, cv2_output=False, **kwargs):
        if extra_fields is None:
            extra_fields = []
        self.transforms = transforms
        self.is_load_label = is_load_label
        self.relation_on = relation_on
        super(VGTSVDataset, self).__init__(yaml_file, cv2_output=cv2_output)

        ignore_attrs = self.cfg.get("ignore_attrs", None)
        # construct those maps
        jsondict_file = find_file_path_in_yaml(self.cfg.get("jsondict", None), self.root)
        jsondict = json.load(open(jsondict_file, 'r'))

        # self.linelist_file
        if 'train' in op.basename(self.linelist_file):
            self.split = "train"
        elif 'test' in op.basename(self.linelist_file) \
                or 'val' in op.basename(self.linelist_file) \
                or 'valid' in op.basename(self.linelist_file):
            self.split = "test"
        else:
            raise ValueError("Split must be one of [train, test], but get {}!".format(self.linelist_file))
        self.filter_duplicate_rels = filter_duplicate_rels and self.split == 'train'

        self.class_to_ind = jsondict['label_to_idx']
        self.ind_to_class = jsondict['idx_to_label']
        self.class_to_ind['__background__'] = 0
        self.ind_to_class['0'] = '__background__'
        self.classes = sort_key_by_val(self.class_to_ind)
        assert (all([self.classes[i] == self.ind_to_class[str(i)] for i in range(len(self.classes))]))

        # writing obj classes to disk for Neural Motif model building.
        obj_classes_out_fn = op.splitext(self.label_file)[0] + ".obj_classes.txt"
        if not op.isfile(obj_classes_out_fn):
            with open(obj_classes_out_fn, 'w') as f:
                for item in self.classes:
                    f.write("%s\n" % item)

        self.attribute_to_ind = jsondict['attribute_to_idx']
        self.ind_to_attribute = jsondict['idx_to_attribute']
        self.attribute_to_ind['__no_attribute__'] = 0
        self.ind_to_attribute['0'] = '__no_attribute__'
        self.attributes = sort_key_by_val(self.attribute_to_ind)
        assert (all([self.attributes[i] == self.ind_to_attribute[str(i)] for i in range(len(self.attributes))]))

        self.relation_to_ind = jsondict['predicate_to_idx']
        self.ind_to_relation = jsondict['idx_to_predicate']
        self.relation_to_ind['__no_relation__'] = 0
        self.ind_to_relation['0'] = '__no_relation__'
        self.relations = sort_key_by_val(self.relation_to_ind)
        assert (all([self.relations[i] == self.ind_to_relation[str(i)] for i in range(len(self.relations))]))

        # writing rel classes to disk for Neural Motif Model building.
        rel_classes_out_fn = op.splitext(self.label_file)[0] + '.rel_classes.txt'
        if not op.isfile(rel_classes_out_fn):
            with open(rel_classes_out_fn, 'w') as f:
                for item in self.relations:
                    f.write("%s\n" % item)

        # label map: minus one because we will add one in BoxLabelLoader
        self.labelmap = {key: val - 1 for key, val in self.class_to_ind.items()}
        labelmap_file = find_file_path_in_yaml(self.cfg.get("labelmap_dec"), self.root)
        # self.labelmap_dec = load_labelmap_file(labelmap_file)
        if self.is_load_label:
            self.label_loader = BoxLabelLoader(
                labelmap=self.labelmap,
                extra_fields=extra_fields,
                ignore_attrs=ignore_attrs
            )

        # get frequency prior for relations
        if self.relation_on:
            self.freq_prior_file = op.splitext(self.label_file)[0] + ".freq_prior.npy"
            if self.split == 'train' and not op.exists(self.freq_prior_file):
                print("Computing frequency prior matrix...")
                fg_matrix, bg_matrix = self._get_freq_prior()
                prob_matrix = fg_matrix.astype(np.float32)
                prob_matrix[:, :, 0] = bg_matrix
                prob_matrix[:, :, 0] += 1
                prob_matrix /= np.sum(prob_matrix, 2)[:, :, None]
                np.save(self.freq_prior_file, prob_matrix)

    def _get_freq_prior(self, must_overlap=False):
        fg_matrix = np.zeros((
            len(self.classes),
            len(self.classes),
            len(self.relations)
        ), dtype=np.int64)

        bg_matrix = np.zeros((
            len(self.classes),
            len(self.classes),
        ), dtype=np.int64)

        for ex_ind in range(self.__len__()):
            target = self.get_groundtruth(ex_ind)
            gt_classes = target.get_field('labels').numpy()
            gt_relations = target.get_field('relation_labels').numpy()
            gt_boxes = target.bbox

            # For the foreground, we'll just look at everything
            try:
                o1o2 = gt_classes[gt_relations[:, :2]]
                for (o1, o2), gtr in zip(o1o2, gt_relations[:, 2]):
                    fg_matrix[o1, o2, gtr] += 1

                # For the background, get all of the things that overlap.
                o1o2_total = gt_classes[np.array(
                    _box_filter(gt_boxes, must_overlap=must_overlap), dtype=int)]
                for (o1, o2) in o1o2_total:
                    bg_matrix[o1, o2] += 1
            except IndexError as e:
                assert len(gt_relations) == 0

            if ex_ind % 20 == 0:
                print("processing {}/{}".format(ex_ind, self.__len__()))

        return fg_matrix, bg_matrix

    def relation_loader(self, relation_triplets, target):
        # relation_triplets [list of tuples]: M*3
        # target: BoxList from label_loader
        if self.filter_duplicate_rels:
            # Filter out dupes!
            assert self.split == 'train'
            all_rel_sets = collections.defaultdict(list)
            for (o0, o1, r) in relation_triplets:
                all_rel_sets[(o0, o1)].append(r)
            relation_triplets = [(k[0], k[1], np.random.choice(v)) for k, v in all_rel_sets.items()]

        # get M*M pred_labels
        relations = torch.zeros([len(target), len(target)], dtype=torch.int64)
        for i in range(len(relation_triplets)):
            subj_id = relation_triplets[i][0]
            obj_id = relation_triplets[i][1]
            pred = relation_triplets[i][2]
            relations[subj_id, obj_id] = int(pred)

        relation_triplets = torch.tensor(relation_triplets)
        target.add_field("relation_labels", relation_triplets)
        target.add_field("pred_labels", relations)
        return target

    def get_target_from_annotations(self, annotations, img_size, idx):
        if self.is_load_label and annotations:
            target = self.label_loader(annotations['objects'], img_size)
            # make sure no boxes are removed
            assert (len(annotations['objects']) == len(target))
            if self.split in ["val", "test"]:
                # add the difficult field
                target.add_field("difficult", torch.zeros(len(target), dtype=torch.int32))
            # load relations
            if self.relation_on:
                target = self.relation_loader(annotations["relations"], target)
            return target

    def get_groundtruth(self, idx, call=False):
        # similar to __getitem__ but without transform
        img = self.get_image(idx)
        if self.cv2_output:
            img_size = img.shape[:2][::-1]  # h, w -> w, h
        else:
            img_size = img.size  # w, h
        annotations = self.get_annotations(idx)
        target = self.get_target_from_annotations(annotations, img_size, idx)
        if call:
            return img, target, annotations
        else:
            return target

    def apply_transforms(self, img, target=None):
        if self.transforms is not None:
            img, target = self.transforms(img, target)
        return img, target

    def map_class_id_to_class_name(self, class_id):
        return self.classes[class_id]

    def map_attribute_id_to_attribute_name(self, attribute_id):
        return self.attributes[attribute_id]

    def map_relation_id_to_relation_name(self, relation_id):
        return self.relations[relation_id]