Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

dapricot test set #1096

Merged
merged 2 commits into from
Jun 8, 2021
Merged

dapricot test set #1096

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3821c48
cherry-picked dapricot test commits from 1088
Jun 8, 2021
41bfa5b
correct checksum filename
Jun 8, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
357 changes: 357 additions & 0 deletions armory/data/adversarial/dapricot_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,357 @@
"""dapricot_test dataset."""

import collections
import json
import os
import pandas

import tensorflow.compat.v1 as tf
import tensorflow_datasets as tfds

_DESCRIPTION = """
LEGAL
-----
Copyright 2021 The MITRE Corporation. All rights reserved.
"""

_CITATION = """
Dataset is unpublished at this time.
"""

_URLS = "https://armory-public-data.s3.us-east-2.amazonaws.com/adversarial-datasets/dapricot_test.tar.gz"


class DapricotTest(tfds.core.GeneratorBasedBuilder):
"""DatasetBuilder for dapricot_test dataset."""

VERSION = tfds.core.Version("1.0.0")
RELEASE_NOTES = {
"1.0.0": "Initial release.",
}

def _info(self) -> tfds.core.DatasetInfo:
"""Returns the dataset metadata."""
features = {
# all Sequences are for [camera_1, camera_2, camera_3]
"image": tfds.features.Sequence(
tfds.features.Image(shape=(None, None, 3)), # encoding_format="jpeg"),
length=3,
),
"images": tfds.features.Sequence(
tfds.features.FeaturesDict(
{
"file_name": tfds.features.Text(),
"height": tf.int64,
"width": tf.int64,
"id": tf.int64,
}
),
length=3,
),
"categories": tfds.features.Sequence(
tfds.features.Sequence(
tfds.features.FeaturesDict(
{
"id": tf.int64, # {'octagon':12, 'diamond':26, 'rect':29}
"name": tfds.features.Text(),
}
)
),
length=3,
),
"objects": tfds.features.Sequence(
tfds.features.Sequence(
{
"id": tf.int64,
"image_id": tf.int64,
"area": tf.int64, # un-normalized area
"boxes": tfds.features.BBoxFeature(), # normalized bounding box [ymin, xmin, ymax, xmax]
"labels": tfds.features.ClassLabel(num_classes=91),
"is_crowd": tf.bool,
}
),
length=3,
),
"patch_metadata": tfds.features.Sequence(
# these data only apply to the "green screen patch" objects
tfds.features.FeaturesDict(
{
"gs_coords": tfds.features.Sequence(
tfds.features.Tensor(
shape=[2], dtype=tf.int64
), # green screen vertices in (x,y)
),
"cc_ground_truth": tfds.features.Tensor(
shape=[24, 3], dtype=tf.float32
), # colorchecker color ground truth
"cc_scene": tfds.features.Tensor(
shape=[24, 3], dtype=tf.float32
), # colorchecker colors in a scene
"shape": tfds.features.Text(), # "diamond", "rect", "octagon"
}
),
length=3,
),
}

return tfds.core.DatasetInfo(
builder=self,
description=_DESCRIPTION,
features=tfds.features.FeaturesDict(features),
citation=_CITATION,
)

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
"""Returns SplitGenerators."""
paths = dl_manager.download_and_extract(_URLS)
return [
tfds.core.SplitGenerator(
name=patch_size,
gen_kwargs={"path": os.path.join(paths, "test"), "size": patch_size},
)
for patch_size in ["large", "medium", "small"]
]

def _generate_examples(self, path, size):
"""yield examples"""

scenes = [
"02",
"03",
"04",
"05",
"07",
"08",
"09",
"10",
"11",
"12",
"13",
"15",
]

size_dist = {"small": "dist15", "medium": "dist10", "large": "dist5"}

yield_id = 0
# For each scene, read JSONs for all cameras.
# For each camera, go through each image.
# For each image, gets its annotations and yield relevant data
for scene in scenes:

annotation_path_camera_1 = os.path.join(
path, "annotations/labels_scene_{}_camera_1.json".format(scene)
)
annotation_path_camera_2 = os.path.join(
path, "annotations/labels_scene_{}_camera_2.json".format(scene)
)
annotation_path_camera_3 = os.path.join(
path, "annotations/labels_scene_{}_camera_3.json".format(scene)
)

dapricot_camera_1 = DapricotAnnotation(annotation_path_camera_1)
dapricot_camera_2 = DapricotAnnotation(annotation_path_camera_2)
dapricot_camera_3 = DapricotAnnotation(annotation_path_camera_3)

images_camera_1 = dapricot_camera_1.images()
images_camera_2 = dapricot_camera_2.images()
images_camera_3 = dapricot_camera_3.images()

# sort images alphabetically so all three cameras are consistent
images_camera_1 = sorted(
images_camera_1, key=lambda x: x["file_name"].lower()
)
images_camera_2 = sorted(
images_camera_2, key=lambda x: x["file_name"].lower()
)
images_camera_3 = sorted(
images_camera_3, key=lambda x: x["file_name"].lower()
)

for image_camera_1, image_camera_2, image_camera_3 in zip(
images_camera_1, images_camera_2, images_camera_3
):

# verify consistency
fname1 = image_camera_1[
"file_name"
] # fname has format "scene_#_camera_1_<SHAPE>_<HEIGHT>_<DIST>.JPG"
fname2 = image_camera_2["file_name"]
fname3 = image_camera_3["file_name"]
assert fname1 == ("_").join(
fname2.split("_")[:3] + ["1"] + fname2.split("_")[4:]
), "{} and {} are inconsistent".format(fname1, fname2)
assert fname1 == ("_").join(
fname3.split("_")[:3] + ["1"] + fname3.split("_")[4:]
), "{} and {} are inconsistent".format(fname1, fname3)

# get object annotations for each image
annotations_camera_1 = dapricot_camera_1.get_annotations(
image_camera_1["id"]
)
annotations_camera_2 = dapricot_camera_2.get_annotations(
image_camera_2["id"]
)
annotations_camera_3 = dapricot_camera_3.get_annotations(
image_camera_3["id"]
)

# convert bbox to Pytorch format
def build_bbox(x, y, width, height):
return tfds.features.BBox(
ymin=y
/ image_camera_1[
"height"
], # all images are the same size, so using image_camera_1 is fine
xmin=x / image_camera_1["width"],
ymax=(y + height) / image_camera_1["height"],
xmax=(x + width) / image_camera_1["width"],
)

# convert segmentation format of (x0,y0,x1,y1,...) to ( (x0, y0), (x1, y1), ... )
def build_coords(segmentation):
xs = segmentation[::2]
ys = segmentation[1::2]
coords = [[int(round(x)), int(round(y))] for (x, y) in zip(xs, ys)]

return coords

# convert green screen shape given in file name to shape expected in downstream algorithms
def get_shape(in_shape):
out_shape = {"stp": "octagon", "pxg": "diamond", "spd": "rect"}
return out_shape[in_shape]

# get colorchecker color box values. There are 24 color boxes, so output shape is (24, 3)
def get_cc(ground_truth=True, scene=None, camera=None):
if ground_truth:
return (
pandas.read_csv(
os.path.join(
path,
"annotations",
"xrite_passport_colors_sRGB-GMB-2005.csv",
),
header=None,
)
.to_numpy()
.astype("float32")
)
else:
return (
pandas.read_csv(
os.path.join(
path,
"annotations",
"scene_{}_camera_{}_CC_values.csv".format(
scene, camera
),
),
header=None,
)
.to_numpy()
.astype("float32")
)

example = {
"image": [
os.path.join(
path,
"scene_{}/camera_{}".format(scene, camera + 1),
im_cam["file_name"],
)
for camera, im_cam in enumerate(
[image_camera_1, image_camera_2, image_camera_3]
)
],
"images": [image_camera_1, image_camera_2, image_camera_3],
"categories": [
d_cam.categories()
for d_cam in [
dapricot_camera_1,
dapricot_camera_2,
dapricot_camera_3,
]
],
"objects": [
[
{
"id": anno["id"],
"image_id": anno["image_id"],
"area": anno["area"],
"boxes": build_bbox(*anno["bbox"]),
"labels": anno["category_id"],
"is_crowd": bool(anno["iscrowd"]),
}
for anno in annos
]
for annos in [
annotations_camera_1,
annotations_camera_2,
annotations_camera_3,
]
],
"patch_metadata": [
[
{
"gs_coords": build_coords(*anno["segmentation"]),
"cc_ground_truth": get_cc(),
"cc_scene": get_cc(
ground_truth=False, scene=scene, camera=camera + 1
),
"shape": get_shape(
im_info["file_name"].split("_")[4].lower()
), # file_name has format "scene_#_camera_#_<SHAPE>_<HEIGHT>_<DIST>.JPG"
}
for anno in annos
if len(anno["segmentation"]) > 0
][0]
for camera, (annos, im_info) in enumerate(
zip(
[
annotations_camera_1,
annotations_camera_2,
annotations_camera_3,
],
[image_camera_1, image_camera_2, image_camera_3],
)
)
],
}

yield_id = yield_id + 1

patch_size = image_camera_1["file_name"].split(".")[
0
] # scene_#_camera_#_<SHAPE>_<HEIGHT>_<DIST>
patch_size = patch_size.split("_")[-1].lower() # <DIST>
if size_dist[size] == patch_size:
yield yield_id, example


class DapricotAnnotation(object):
"""Dapricot annotation helper class."""

def __init__(self, annotation_path):
with tf.io.gfile.GFile(annotation_path) as f:
data = json.load(f)
self._data = data

# for each images["id"], find all annotations such that annotations["image_id"] == images["id"]
img_id2annotations = collections.defaultdict(list)
for a in self._data["annotations"]:
img_id2annotations[a["image_id"]].append(a)
self._img_id2annotations = {
k: list(sorted(v, key=lambda a: a["id"]))
for k, v in img_id2annotations.items()
}

def categories(self):
"""Return the category dicts, as sorted in the file."""
return self._data["categories"]

def images(self):
"""Return the image dicts, as sorted in the file."""
return self._data["images"]

def get_annotations(self, img_id):
"""Return all annotations associated with the image id string."""
return self._img_id2annotations.get(img_id, [])
Loading