def __init__(self, mode):
"""
:param mode: train2014 or val2014
"""
self.mode = mode
self.root = os.path.join(COCO_PATH, mode)
self.ann_file = os.path.join(COCO_PATH, 'annotations', 'instances_{}.json'.format(mode))
self.coco = COCO(self.ann_file)
self.ids = [k for k in self.coco.imgs.keys() if len(self.coco.imgToAnns[k]) > 0]
tform = []
if self.is_train:
tform.append(RandomOrder([
Grayscale(),
Brightness(),
Contrast(),
Sharpness(),
Hue(),
]))
tform += [
SquarePad(),
Resize(IM_SCALE),
ToTensor(),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
self.transform_pipeline = Compose(tform)
self.ind_to_classes = ['__background__'] + [v['name'] for k, v in self.coco.cats.items()]
# COCO inds are weird (84 inds in total but a bunch of numbers are skipped)
self.id_to_ind = {coco_id:(ind+1) for ind, coco_id in enumerate(self.coco.cats.keys())}
self.id_to_ind[0] = 0
self.ind_to_id = {x:y for y,x in self.id_to_ind.items()}
def __init__(self, mode, filter_duplicate_rels=True, mask_resolution=28, use_for_bias=False):
"""
:param mode: train2014 or val2014
"""
self.mask_resolution = mask_resolution
self.mode = mode
self.filter_duplicate_rels = filter_duplicate_rels and self.mode == 'train'
tform = []
if self.is_train:
tform.append(RandomOrder([
Grayscale(),
Brightness(),
Contrast(),
Sharpness(),
Hue(),
]))
tform += [
SquarePad(),
Resize(IM_SCALE),
ToTensor(),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
self.transform_pipeline = Compose(tform)
#image_names = [name[:-4] for name in os.listdir(os.path.join(PIC_PATH, 'image/'+mode)) if name.endswith('.jpg')]
error_list = [line.strip()[:-4] for line in open(os.path.join(PIC_PATH, mode+'_error_list.txt'))]
image_names = [name[:-4] for name in os.listdir(os.path.join(PIC_PATH, 'image/' + mode)) if
name.endswith('.jpg') and name[:-4] not in error_list]
image_names.sort(key=str.lower)
if self.mode != 'test':
semantic_names = [name[:-4] for name in os.listdir(os.path.join(PIC_PATH, 'segmentation/'+mode+'/semantic'))
if name.endswith('.png') and name[:-4] not in error_list]
instance_names = [name[:-4] for name in os.listdir(os.path.join(PIC_PATH, 'segmentation/'+mode+'/instance'))
if name.endswith('.png') and name[:-4] not in error_list]
semantic_names.sort(key=str.lower)
instance_names.sort(key=str.lower)
assert image_names == semantic_names
assert image_names == instance_names
# image_names = [name[:-4] for name in os.listdir(os.path.join(PIC_OFFLINE_PATH, 'val/obj_feat/')) if
# name.endswith('.npy')]
# image_names.sort(key=str.lower)
self.img_names = image_names
print(len(self.img_names))
rel_cats = json.load(open(os.path.join(PIC_PATH,'categories_list/relation_categories.json')))
self.ind_to_predicates = [rel_cat['name'] for rel_cat in rel_cats]
cls_cats = json.load(open(os.path.join(PIC_PATH, 'categories_list/label_categories.json')))
self.ind_to_classes = [cls_cat['name'] for cls_cat in cls_cats]
if self.mode != 'test':
self.img2rels = dict()
img_relations = json.load(open(os.path.join(PIC_PATH, 'relation/relations_'+self.mode+'.json')))
for img_relation in img_relations:
rels = []
for index, rel in enumerate(img_relation['relations']):
temp = np.array([[rel['subject']-1, rel['object']-1, rel['relation']]], dtype=np.int32)
rels.append(temp)
rels = np.concatenate(rels, axis=0)
self.img2rels[img_relation['name'][:-4]] = rels
print('====================')
print(self.ind_to_classes)
print(self.ind_to_predicates)
self.id_to_ind = {ind: ind for ind, name in enumerate(self.ind_to_classes)}
self.ind_to_id = {x: y for y, x in self.id_to_ind.items()}
#self.create_coco_format()
if use_for_bias:
dataset = json.load(open(os.path.join(PIC_PATH, 'train.json')))
anns, imgs = {}, {}
imgToAnns = defaultdict(list)
for img in dataset['images']:
imgs[img['id']] = img['file_name'][:-4]
for ann in dataset['annotations']:
imgToAnns[imgs[ann['image_id']]].append(ann)
self.img2boxes = dict()
self.img2classes = dict()
for img_name, anns in imgToAnns.items():
gt_box = []
gt_class = []
for ann in anns:
gt_box.append(np.array([ann['bbox']]))
gt_class.append(np.array([ann['category_id']]))
gt_box = np.concatenate(gt_box, axis=0)
gt_class = np.concatenate(gt_class, axis=0)
self.img2boxes[img_name] = gt_box
self.img2classes[img_name] = gt_class
def __init__(
self,
mode,
# image_dir, instances_json, stuff_json=None,
stuff_only=True,
image_size=(64, 64),
mask_size=16,
normalize_images=True,
max_samples=None,
include_relationships=True,
min_object_size=0.02,
min_objects_per_image=3,
max_objects_per_image=8,
include_other=False,
instance_whitelist=None,
stuff_whitelist=None):
"""
A PyTorch Dataset for loading Coco and Coco-Stuff annotations and converting
them to scene graphs on the fly.
Inputs:
- image_dir: Path to a directory where images are held
- instances_json: Path to a JSON file giving COCO annotations
- stuff_json: (optional) Path to a JSON file giving COCO-Stuff annotations
- stuff_only: (optional, default True) If True then only iterate over
images which appear in stuff_json; if False then iterate over all images
in instances_json.
- image_size: Size (H, W) at which to load images. Default (64, 64).
- mask_size: Size M for object segmentation masks; default 16.
- normalize_image: If True then normalize images by subtracting ImageNet
mean pixel and dividing by ImageNet std pixel.
- max_samples: If None use all images. Other wise only use images in the
range [0, max_samples). Default None.
- include_relationships: If True then include spatial relationships; if
False then only include the trivial __in_image__ relationship.
- min_object_size: Ignore objects whose bounding box takes up less than
this fraction of the image.
- min_objects_per_image: Ignore images which have fewer than this many
object annotations.
- max_objects_per_image: Ignore images which have more than this many
object annotations.
- include_other: If True, include COCO-Stuff annotations which have category
"other". Default is False, because I found that these were really noisy
and pretty much impossible for the system to model.
- instance_whitelist: None means use all instance categories. Otherwise a
list giving a whitelist of instance category names to use.
- stuff_whitelist: None means use all stuff categories. Otherwise a list
giving a whitelist of stuff category names to use.
"""
super(Dataset, self).__init__()
self.mode = mode
image_dir = join(COCO_PATH, "images", "%s2017" % mode)
instances_json = join(COCO_PATH, "annotations",
"instances_%s2017.json" % mode)
stuff_json = join(COCO_PATH, "annotations", "stuff_%s2017.json" % mode)
if stuff_only and stuff_json is None:
print('WARNING: Got stuff_only=True but stuff_json=None.')
print('Falling back to stuff_only=False.')
self.image_dir = image_dir
self.mask_size = mask_size
self.max_samples = max_samples
self.normalize_images = normalize_images
self.include_relationships = include_relationships
# self.set_image_size(image_size)
tform = []
if self.is_train:
tform.append(
RandomOrder([
Grayscale(),
Brightness(),
Contrast(),
Sharpness(),
Hue(),
]))
tform += [
SquarePad(),
Resize(IM_SCALE),
ToTensor(),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
self.transform_pipeline = Compose(tform)
with open(instances_json, 'r') as f:
instances_data = json.load(f)
stuff_data = None
if stuff_json is not None and stuff_json != '':
with open(stuff_json, 'r') as f:
stuff_data = json.load(f)
self.image_ids = []
self.image_id_to_filename = {}
self.image_id_to_size = {}
for image_data in instances_data['images']:
image_id = image_data['id']
filename = image_data['file_name']
width = image_data['width']
height = image_data['height']
self.image_ids.append(image_id)
self.image_id_to_filename[image_id] = filename
self.image_id_to_size[image_id] = (width, height)
self.vocab = {
'object_name_to_idx': {},
'pred_name_to_idx': {},
}
object_idx_to_name = {}
all_instance_categories = []
for category_data in instances_data['categories']:
category_id = category_data['id']
category_name = category_data['name']
all_instance_categories.append(category_name)
object_idx_to_name[category_id] = category_name
self.vocab['object_name_to_idx'][category_name] = category_id
all_stuff_categories = []
if stuff_data:
for category_data in stuff_data['categories']:
category_name = category_data['name']
category_id = category_data['id']
all_stuff_categories.append(category_name)
object_idx_to_name[category_id] = category_name
self.vocab['object_name_to_idx'][category_name] = category_id
if instance_whitelist is None:
instance_whitelist = all_instance_categories
if stuff_whitelist is None:
stuff_whitelist = all_stuff_categories
category_whitelist = set(instance_whitelist) | set(stuff_whitelist)
# Add object data from instances
self.image_id_to_objects = defaultdict(list)
for object_data in instances_data['annotations']:
image_id = object_data['image_id']
_, _, w, h = object_data['bbox']
W, H = self.image_id_to_size[image_id]
box_area = (w * h) / (W * H)
box_ok = box_area > min_object_size
object_name = object_idx_to_name[object_data['category_id']]
category_ok = object_name in category_whitelist
other_ok = object_name != 'other' or include_other
if box_ok and category_ok and other_ok:
self.image_id_to_objects[image_id].append(object_data)
# Add object data from stuff
if stuff_data:
image_ids_with_stuff = set()
for object_data in stuff_data['annotations']:
image_id = object_data['image_id']
image_ids_with_stuff.add(image_id)
_, _, w, h = object_data['bbox']
W, H = self.image_id_to_size[image_id]
box_area = (w * h) / (W * H)
box_ok = box_area > min_object_size
object_name = object_idx_to_name[object_data['category_id']]
category_ok = object_name in category_whitelist
other_ok = object_name != 'other' or include_other
if box_ok and category_ok and other_ok:
self.image_id_to_objects[image_id].append(object_data)
if stuff_only:
new_image_ids = []
for image_id in self.image_ids:
if image_id in image_ids_with_stuff:
new_image_ids.append(image_id)
self.image_ids = new_image_ids
all_image_ids = set(self.image_id_to_filename.keys())
image_ids_to_remove = all_image_ids - image_ids_with_stuff
for image_id in image_ids_to_remove:
self.image_id_to_filename.pop(image_id, None)
self.image_id_to_size.pop(image_id, None)
self.image_id_to_objects.pop(image_id, None)
# COCO category labels start at 1, so use 0 for __image__
self.vocab['object_name_to_idx']['__image__'] = 0
# Build object_idx_to_name
name_to_idx = self.vocab['object_name_to_idx']
assert len(name_to_idx) == len(set(name_to_idx.values()))
max_object_idx = max(name_to_idx.values())
idx_to_name = ['NONE'] * (1 + max_object_idx)
for name, idx in self.vocab['object_name_to_idx'].items():
idx_to_name[idx] = name
self.vocab['object_idx_to_name'] = idx_to_name
# Prune images that have too few or too many objects
new_image_ids = []
total_objs = 0
for image_id in self.image_ids:
num_objs = len(self.image_id_to_objects[image_id])
total_objs += num_objs
if min_objects_per_image <= num_objs <= max_objects_per_image:
new_image_ids.append(image_id)
# self.image_ids = new_image_ids
self.ids = new_image_ids
self.vocab['pred_idx_to_name'] = [
'__in_image__',
'left of',
'right of',
'above',
'below',
'inside',
'surrounding',
]
self.vocab['pred_name_to_idx'] = {}
for idx, name in enumerate(self.vocab['pred_idx_to_name']):
self.vocab['pred_name_to_idx'][name] = idx
# for object detection model to get number of class
self.ind_to_classes = self.vocab['object_idx_to_name']
self.ind_to_id = {i: i for i in range(len(self.ind_to_classes))}
def __init__(self, mode, filter_duplicate_rels=True, mask_resolution=28):
"""
:param mode: train2014 or val2014
"""
self.mask_resolution = mask_resolution
self.mode = mode
self.filter_duplicate_rels = filter_duplicate_rels and self.mode == 'train'
tform = []
if self.is_train:
tform.append(
RandomOrder([
Grayscale(),
Brightness(),
Contrast(),
Sharpness(),
Hue(),
]))
tform += [
SquarePad(),
Resize(IM_SCALE),
ToTensor(),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
self.transform_pipeline = Compose(tform)
image_names = [
name[:-4]
for name in os.listdir(os.path.join(PIC_PATH, 'image/' + mode))
if name.endswith('.jpg')
]
if self.mode != 'test':
semantic_names = [
name[:-4] for name in os.listdir(
os.path.join(PIC_PATH, 'segmentation/' + mode +
'/semantic')) if name.endswith('.png')
]
instance_names = [
name[:-4] for name in os.listdir(
os.path.join(PIC_PATH, 'segmentation/' + mode +
'/instance')) if name.endswith('.png')
]
image_names.sort(key=str.lower)
semantic_names.sort(key=str.lower)
instance_names.sort(key=str.lower)
assert image_names == semantic_names
assert image_names == instance_names
# image_names = [name[:-4] for name in os.listdir(os.path.join(PIC_OFFLINE_PATH, 'val/obj_feat/')) if
# name.endswith('.npy')]
# image_names.sort(key=str.lower)
self.img_names = image_names
rel_cats = json.load(
open(
os.path.join(PIC_PATH,
'categories_list/relation_categories.json')))
self.ind_to_predicates = [rel_cat['name'] for rel_cat in rel_cats]
cls_cats = json.load(
open(
os.path.join(PIC_PATH,
'categories_list/label_categories.json')))
self.ind_to_classes = [cls_cat['name'] for cls_cat in cls_cats]
if self.mode != 'test':
self.img2rels = dict()
img_relations = json.load(
open(
os.path.join(PIC_PATH,
'relation/relations_' + self.mode + '.json')))
for img_relation in img_relations:
rels = []
for index, rel in enumerate(img_relation['relations']):
temp = np.array([[
rel['subject'] - 1, rel['object'] - 1, rel['relation']
]],
dtype=np.int32)
rels.append(temp)
rels = np.concatenate(rels, axis=0)
self.img2rels[img_relation['name'][:-4]] = rels
print('====================')
print(self.ind_to_classes)
print(self.ind_to_predicates)
self.id_to_ind = {
ind: ind
for ind, name in enumerate(self.ind_to_classes)
}
self.ind_to_id = {x: y for y, x in self.id_to_ind.items()}