File size: 8,576 Bytes
25986db
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
import json
import os
import random
from collections import OrderedDict

import torch

from lib.train.admin import env_settings
from lib.train.data import jpeg4py_loader
from .base_video_dataset import BaseVideoDataset
from .refer import REFER
class RefCOCOSeq(BaseVideoDataset):
    """ The COCO dataset. COCO is an image dataset. Thus, we treat each image as a sequence of length 1.



    Publication:

        Microsoft COCO: Common Objects in Context.

        Tsung-Yi Lin, Michael Maire, Serge J. Belongie, Lubomir D. Bourdev, Ross B. Girshick, James Hays, Pietro Perona,

        Deva Ramanan, Piotr Dollar and C. Lawrence Zitnick

        ECCV, 2014

        https://arxiv.org/pdf/1405.0312.pdf



    Download the images along with annotations from http://cocodataset.org/#download. The root folder should be

    organized as follows.

        - coco_root

            - annotations

                - instances_train2014.json

                - instances_train2017.json

            - images

                - train2014

                - train2017



    Note: You also have to install the coco pythonAPI from https://github.com/cocodataset/cocoapi.

    """

    def __init__(self, root=None, image_loader=jpeg4py_loader, data_fraction=None, split="train", version="2014",

                 name="refcoco", splitBy="google"):
        """

        args:

            root - path to the coco dataset.

            image_loader (default_image_loader) -  The function to read the images. If installed,

                                                   jpeg4py (https://github.com/ajkxyz/jpeg4py) is used by default. Else,

                                                   opencv's imread is used.

            data_fraction (None) - Fraction of images to be used. The images are selected randomly. If None, all the

                                  images  will be used

            split - 'train' or 'val'.

            version - version of coco dataset (2014 or 2017)

        """
        root = env_settings().coco_dir if root is None else root
        root = os.path.join(root,"other")
        super().__init__('RefCOCOSeq', root, image_loader)
        self.split = split
        self.img_pth = os.path.join(root, 'images/mscoco/images/{}{}'.format("train", version))
        self.anno_path = os.path.join(root, '{}/instances.json'.format(name))
        self.dataset_name = name
        # Load the COCO set.
        self.coco_set = REFER(root, dataset=name, splitBy=splitBy)

        self.cats = self.coco_set.Cats

        self.class_list = self.get_class_list()

        self.sequence_list = self._get_sequence_list()

        if data_fraction is not None:
            self.sequence_list = random.sample(self.sequence_list, int(len(self.sequence_list) * data_fraction))
        self.seq_per_class = self._build_seq_per_class()

        # get_the_subject_infor
        self.subject_infor_path = "./"
        while "resource" not in os.listdir(self.subject_infor_path):
            self.subject_infor_path = os.path.join(self.subject_infor_path, "../")

        self.subject_infor_path = os.path.join(self.subject_infor_path,
                                               "resource/text_infor/refcoco_4o_v3_extract_mask.json")
        with open(self.subject_infor_path, 'r') as file:
            self.subject_infor = json.load(file)

    def _get_sequence_list(self):
        ref_list = list(self.coco_set.getRefIds(split=self.split))
        seq_list = [a for a in ref_list if self.coco_set.refToAnn[a]['iscrowd'] == 0]

        return seq_list

    def is_video_sequence(self):
        return False

    def get_num_classes(self):
        return len(self.class_list)

    def get_name(self):
        return self.dataset_name

    def has_class_info(self):
        return True

    def get_class_list(self):
        class_list = []
        for cat_id in self.cats.keys():
            class_list.append(self.cats[cat_id])
        return class_list

    def has_segmentation_info(self):
        return True

    def get_num_sequences(self):
        return len(self.sequence_list)

    def _build_seq_per_class(self):
        seq_per_class = {}
        for i, seq in enumerate(self.sequence_list):
            class_name = self.cats[self.coco_set.refToAnn[seq]['category_id']]
            if class_name not in seq_per_class:
                seq_per_class[class_name] = [i]
            else:
                seq_per_class[class_name].append(i)

        return seq_per_class

    def get_sequences_in_class(self, class_name):
        return self.seq_per_class[class_name]

    def get_sequence_info(self, seq_id):
        anno = self._get_anno(seq_id)

        bbox = torch.Tensor(anno['bbox']).view(1, 4)

        # mask = torch.Tensor(self.coco_set.annToMask(anno)).unsqueeze(dim=0)

        '''2021.1.3 To avoid too small bounding boxes. Here we change the threshold to 50 pixels'''
        valid = (bbox[:, 2] > 50) & (bbox[:, 3] > 50)

        visible = valid.clone().byte()

        nlp = self._read_nlp(seq_id)

        return {'bbox': bbox, 'valid': valid, 'visible': visible, 'nlp': nlp}

    def _read_nlp(self, seq_id):

        ref = self.coco_set.Refs[self.sequence_list[seq_id]]
        sent = ref['sentences'][-1]['sent']

        return sent

    def _get_anno(self, seq_id):
        anno = self.coco_set.refToAnn[self.sequence_list[seq_id]]

        return anno

    def _get_frames(self, seq_id):
        path = self.coco_set.loadImgs([self.coco_set.refToAnn[self.sequence_list[seq_id]]['image_id']])[0]['file_name']
        img = self.image_loader(os.path.join(self.img_pth, path))
        return img

    def get_meta_info(self, seq_id):
        try:
            cat_dict_current = self.cats[self.coco_set.refToAnn[self.sequence_list[seq_id]]['category_id']]
            object_meta = OrderedDict({'object_class_name': cat_dict_current['name'],
                                       'motion_class': None,
                                       'major_class': cat_dict_current['supercategory'],
                                       'root_class': None,
                                       'motion_adverb': None})
        except:
            object_meta = OrderedDict({'object_class_name': None,
                                       'motion_class': None,
                                       'major_class': None,
                                       'root_class': None,
                                       'motion_adverb': None})
        return object_meta

    def get_class_name(self, seq_id):
        cat_dict_current = self.cats[self.coco_set.refToAnn[self.sequence_list[seq_id]]['category_id']]
        return cat_dict_current['name']
    def get_subject_mask_infor(self,seq_id):
        # seq_name = self.sequence_list[seq_id]
        seq_sub_mask_infor = self.subject_infor[str(seq_id)]["subject_extrack_mask_infor"]

        return seq_sub_mask_infor

    def get_frames(self, seq_id=None, frame_ids=None, anno=None):
        # COCO is an image dataset. Thus we replicate the image denoted by seq_id len(frame_ids) times, and return a
        # list containing these replicated images.
        frame = self._get_frames(seq_id)

        seq_name = self.sequence_list[seq_id]

        frame_list = [frame.copy() for _ in frame_ids]

        if anno is None:
            anno = self.get_sequence_info(seq_id)

        anno_frames = {}
        for key, value in anno.items():
            if key == 'nlp':
                anno_frames[key] = [value for _ in frame_ids]
            else:
                anno_frames[key] = [value[0, ...] for _ in frame_ids]

        object_meta = self.get_meta_info(seq_id)
        subject_mask_infor = self.get_subject_mask_infor(seq_id)
        if "NoneNone" in subject_mask_infor:
            subject_mask_infor = [-1]
        nlp_with_mask = anno_frames["nlp"][0] + "+" + ",".join(map(str, subject_mask_infor))

        anno_frames["nlp"] = [nlp_with_mask]

        # anno_frames["subject_mask_infor"] = subject_mask_infor
        return frame_list, anno_frames, object_meta

    def get_path(self, seq_id, frame_ids):
        img_name = self.coco_set.loadImgs([self.coco_set.refToAnn[self.sequence_list[seq_id]]['image_id']])[0]['file_name']
        return [os.path.join(self.img_pth, img_name) for _ in range(len(frame_ids))]

    def get_ref_id(self, seq_id):
        return self.sequence_list[seq_id]