def _get_data(stem: str, image_path: str, annotation_path: str) -> Data:
try:
import xmltodict # pylint: disable=import-outside-toplevel
except ModuleNotFoundError as error:
raise ModuleImportError(module_name=error.name) from error
data = Data(os.path.join(image_path, f"{stem}.jpg"))
box2d = []
with open(os.path.join(annotation_path, f"{stem}.xml"),
encoding="utf-8") as fp:
labels: Any = xmltodict.parse(fp.read())
objects = labels["annotation"]["object"]
if not isinstance(objects, list):
objects = [objects]
for item in objects:
category = item["name"]
attributes = {k: bool(int(v)) for k, v in item["actions"].items()}
bndbox = item["bndbox"]
box2d.append(
LabeledBox2D(
float(bndbox["xmin"]),
float(bndbox["ymin"]),
float(bndbox["xmax"]),
float(bndbox["ymax"]),
category=category,
attributes=attributes,
))
data.label.box2d = box2d
return data
def _get_data(mat: Any, name: Any, bbox: Any, file_path: str) -> Data:
image_path = "".join(chr(v[0]) for v in mat[name[0]])
data = Data(os.path.join(file_path, image_path),
target_remote_path=image_path.zfill(10))
data.label.box2d = []
mat_bbox = mat[bbox[0]]
labeled_box = ({key: [value[0][0]]
for key, value in mat_bbox.items()}
if mat_bbox["label"].shape[0] == 1 else {
key: [mat[value[0]][0][0] for value in values]
for key, values in mat_bbox.items()
})
for x, y, w, h, e in zip(
labeled_box["left"],
labeled_box["top"],
labeled_box["width"],
labeled_box["height"],
labeled_box["label"],
):
data.label.box2d.append(
LabeledBox2D.from_xywh(x,
y,
w,
h,
category="0" if e == 10 else str(int(e))))
return data
def _parse_xml(xml_path: str) -> Dict[int, List[LabeledBox2D]]:
dom = parse(xml_path)
frames = dom.documentElement.getElementsByTagName("frame")
labels = defaultdict(list)
for frame in frames:
index = int(frame.getAttribute("number"))
objects = frame.getElementsByTagName("objectlist")[0]
for obj in objects.getElementsByTagName("object"):
box = obj.getElementsByTagName("box")[0]
box_h = int(box.getAttribute("h"))
box_w = int(box.getAttribute("w"))
box_xc = int(box.getAttribute("xc"))
box_yc = int(box.getAttribute("yc"))
labels[index].append(
LabeledBox2D.from_xywh(
x=box_xc - box_w // 2,
y=box_yc - box_h // 2,
width=box_w,
height=box_h,
category=obj.getElementsByTagName("subtype")
[0].firstChild.data,
))
return labels
def _extract_box2d(
path: str, sequence: str,
box2d_catalog: Box2DSubcatalog) -> Dict[int, List[LabeledBox2D]]:
attributes = box2d_catalog.attributes
category_names = box2d_catalog.categories.keys()
out_of_view_level = attributes["out_of_view"].enum
occlusion_level = attributes["occlusion"].enum
ground_truth_path = os.path.join(path, "UAV-benchmark-MOTD_v1.0", "GT")
frame_id_ground_truth_map = defaultdict(list)
with open(os.path.join(ground_truth_path, f"{sequence}_gt_whole.txt"),
encoding="utf-8") as fp:
for elements in csv.DictReader(fp, fieldnames=_FIELDNAMES):
box2d = LabeledBox2D.from_xywh(
*(int(elements[key]) for key in _XYWH_KEYS),
category=category_names[int(elements["object_category"]) - 1],
attributes={
"out_of_view":
out_of_view_level[int(elements["out_of_view"]) - 1],
"occlusion":
occlusion_level[int(elements["occlusion"]) - 1],
},
instance=elements["target_id"],
)
frame_id = int(elements["frame_index"])
frame_id_ground_truth_map[frame_id].append(box2d)
return frame_id_ground_truth_map
def _create_box_label(
label_path: str,
box2d_subcatalog: Box2DSubcatalog,
model_to_attributes: Dict[str, Dict[str, Union[int, float, str, None]]],
) -> List[LabeledBox2D]:
# label_path looks like:
# <root_path>/<make_name_id>/<model_name_id>/<release_year>/<file_name>.txt
_, make_id, model_id, release_year, _ = label_path.rsplit(os.sep, 4)
with open(label_path, encoding="utf-8") as fp:
viewpoint_id = int(fp.readline().strip())
viewpoint_index = viewpoint_id if viewpoint_id == -1 else viewpoint_id - 1
attributes = model_to_attributes[model_id].copy()
attributes["released_year"] = int(
release_year) if release_year != "unknown" else None
attributes["viewpoint"] = box2d_subcatalog.attributes[
"viewpoint"].enum[viewpoint_index]
fp.readline()
category = box2d_subcatalog.categories[int(make_id) - 1].name
box2d_label = [
LabeledBox2D(*map(int,
line.strip().split()),
category=category,
attributes=attributes) for line in fp
]
return box2d_label
def _get_instance_label(instances_annotations: Dict[int, Any], image_id: int,
categories: Dict[int, str]) -> Label:
label: Label = Label()
label.box2d = []
label.multi_polygon = []
label.rle = []
if image_id not in instances_annotations:
return label
for annotation in instances_annotations[image_id]:
category = categories[annotation["category_id"]]
label.box2d.append(
LabeledBox2D.from_xywh(*annotation["bbox"], category=category))
if annotation["iscrowd"] == 0:
points = [
chunked(coordinates, 2)
for coordinates in annotation["segmentation"]
]
label.multi_polygon.append(
LabeledMultiPolygon(points, category=category))
else:
label.rle.append(
LabeledRLE(annotation["segmentation"]["counts"],
category=category))
return label
def _get_data(filename: str, image_path: str, annotation_path: str) -> Data:
"""Get all information of the datum corresponding to filename.
Arguments:
filename: The filename of the data.
image_path: The path of the image directory.
annotation_path: The path of the annotation directory.
Returns:
Data: class: `~tensorbay.dataset.data.Data` instance.
"""
data = Data(os.path.join(image_path, f"{filename}.jpg"))
box2d = []
with open(os.path.join(annotation_path, f"{filename}.xml"), "r", encoding="utf-8") as fp:
objects = xmltodict.parse(fp.read())["annotation"]["object"]
if not isinstance(objects, list):
objects = [objects]
for obj in objects:
attributes = {attribute: bool(int(obj[attribute])) for attribute in _BOOLEAN_ATTRIBUTES}
attributes["pose"] = obj["pose"]
bndbox = obj["bndbox"]
box2d.append(
LabeledBox2D(
float(bndbox["xmin"]),
float(bndbox["ymin"]),
float(bndbox["xmax"]),
float(bndbox["ymax"]),
category=obj["name"],
attributes=attributes,
)
)
data.label.box2d = box2d
return data
def AnimalPose7(path: str) -> Dataset:
"""`7 Categories Animal-Pose <https://sites.google.com/view/animal-pose/>`_ dataset.
The file structure should be like::
<path>
bndbox_image/
antelope/
Img-77.jpg
...
...
bndbox_anno/
antelope.json
...
Arguments:
path: The root directory of the dataset.
Returns:
loaded :class:`~tensorbay.dataset.dataset.Dataset` object.
"""
root_path = os.path.abspath(os.path.expanduser(path))
dataset = Dataset(DATASET_NAME_7)
dataset.load_catalog(
os.path.join(os.path.dirname(__file__), "catalog_7.json"))
segment = dataset.create_segment()
for animal in dataset.catalog.box2d.categories.keys():
with open(os.path.join(root_path, "bndbox_anno", f"{animal}.json"),
encoding="utf-8") as fp:
annotations = json.load(fp)
for image_name, box2ds in annotations.items():
image_path = os.path.join(root_path, "bndbox_image", animal,
image_name)
data = Data(image_path,
target_remote_path=f"{animal}/{image_name}")
data.label.box2d = []
for box2d in box2ds:
coordinates = box2d["bndbox"]
data.label.box2d.append(
LabeledBox2D(
float(coordinates["xmin"]),
float(coordinates["ymin"]),
float(coordinates["xmax"]),
float(coordinates["ymax"]),
category=animal,
))
segment.append(data)
return dataset
def _add_box2d_label(label_info: Dict[str, Any], box2d: List[LabeledBox2D]) -> None:
box2d_info = label_info["box2d"]
labeled_box2d = LabeledBox2D(
box2d_info["x1"],
box2d_info["y1"],
box2d_info["x2"],
box2d_info["y2"],
category=label_info["category"],
attributes=label_info["attributes"],
)
box2d.append(labeled_box2d)
def test_eq(self):
box2d1 = LabeledBox2D(1,
1,
3,
3,
category="cat",
attributes={"gender": "male"})
box2d2 = LabeledBox2D(1,
1,
3,
3,
category="cat",
attributes={"gender": "male"})
box2d3 = LabeledBox2D(1,
1,
4,
4,
category="cat",
attributes={"gender": "male"})
assert box2d1 == box2d2
assert box2d1 != box2d3
def _get_data_part1(root_path: str, aniamls: Iterable[str]) -> Iterator[Data]:
try:
import xmltodict # pylint: disable=import-outside-toplevel
except ModuleNotFoundError as error:
raise ModuleImportError(module_name=error.name) from error
for animal in aniamls:
for image_path in glob(
os.path.join(root_path, "keypoint_image_part1", animal,
"*.jpg")):
data = Data(
image_path,
target_remote_path=f"{animal}/{os.path.basename(image_path)}")
for annotation_path in glob(
os.path.join(
root_path,
"PASCAL2011_animal_annotation",
animal,
f"{os.path.splitext(os.path.basename(image_path))[0]}_*.xml",
)):
with open(annotation_path, encoding="utf-8") as fp:
labels: Any = xmltodict.parse(fp.read())
box2d = labels["annotation"]["visible_bounds"]
data.label.box2d = [
LabeledBox2D.from_xywh(
x=float(box2d["@xmin"]),
y=float(box2d["@ymin"]),
width=float(box2d["@width"]),
height=float(box2d["@height"]),
category=animal,
)
]
keypoints2d: List[Tuple[float, float, int]] = [
()
] * 20 # type: ignore[list-item]
for keypoint in labels["annotation"]["keypoints"]["keypoint"]:
keypoints2d[_KEYPOINT_TO_INDEX[keypoint["@name"]]] = (
float(keypoint["@x"]),
float(keypoint["@y"]),
int(keypoint["@visible"]),
)
data.label.keypoints2d = [
LabeledKeypoints2D(keypoints2d, category=animal)
]
yield data
def _get_object_annotations(
object_annotations: List[Dict[str, Any]],
token_to_categories: Dict[str, Any],
token_to_attributes: Dict[str, Any],
) -> Iterator[Tuple[LabeledBox2D, Optional[LabeledRLE]]]:
for object_annotation in object_annotations:
category = token_to_categories[object_annotation["category_token"]]["name"]
attributes = {}
for attribute_token in object_annotation["attribute_tokens"]:
key, value = token_to_attributes[attribute_token]["name"].rsplit(".", 1)
attributes[key] = value
box2d = LabeledBox2D(*object_annotation["bbox"], category=category, attributes=attributes)
mask = object_annotation["mask"]
rle = _get_labeled_rle(mask, category, attributes) if mask else None
yield box2d, rle
def _get_box_label(file_path: str) -> List[LabeledBox2D]:
with open(file_path, "r", encoding="utf-8") as fp:
objects = xmltodict.parse(fp.read())["annotation"]["object"]
if not isinstance(objects, list):
objects = [objects]
box2d = []
for obj in objects:
bndbox = obj["bndbox"]
box2d.append(
LabeledBox2D(
float(bndbox["xmin"]),
float(bndbox["ymin"]),
float(bndbox["xmax"]),
float(bndbox["ymax"]),
))
return box2d
def _get_panoptic_box2d(panoptic_annotations: Dict[int, Any], image_id: int,
categories: Dict[int, str]) -> List[LabeledBox2D]:
if image_id not in panoptic_annotations:
return []
box2d: List[LabeledBox2D] = []
for annotation in panoptic_annotations[image_id]:
for segment_info in annotation["segments_info"]:
# category_id 1-91 are thing categories from the detection task
# category_id 92-200 are stuff categories from the stuff task
if segment_info["category_id"] > 91:
category = categories[segment_info["category_id"]]
box2d.append(
LabeledBox2D.from_xywh(*segment_info["bbox"],
category=category))
return box2d
def test_dumps(self):
contents = {
"CLASSIFICATION": {"category": "cat", "attributes": {"gender": "male"}},
"BOX2D": [
{
"box2d": {"xmin": 1, "ymin": 1, "xmax": 2, "ymax": 2},
"category": "dog",
"attributes": {"gender": "female"},
}
],
}
label = Label()
label.classification = Classification.loads(contents["CLASSIFICATION"])
label.box2d = [LabeledBox2D.loads(contents["BOX2D"][0])]
assert label.dumps() == contents
def _load_box_labels(file_path: str) -> List[LabeledBox2D]:
box_labels = []
with open(file_path, encoding="utf-8") as fp:
for line in fp:
center_x, center_y, width, height, occlusion, blur = map(int, line.strip().split())
attributes = {"occlusion-level": _OCCLUSION_MAP[occlusion], "blur-level": bool(blur)}
box_labels.append(
LabeledBox2D.from_xywh(
x=center_x - width / 2,
y=center_y - height / 2,
width=width,
height=height,
attributes=attributes,
)
)
return box_labels
def _add_labels(segment: Segment, csv_path: str) -> None:
supercategory = _get_supercategory(csv_path)
with open(csv_path, "r", encoding="utf-8") as fp:
reader = csv.DictReader(fp, delimiter=";")
for row in reader:
frame_number = int(row["Origin frame number"])
data = segment[frame_number]
label = LabeledBox2D(
int(row["Upper left corner X"]),
int(row["Upper left corner Y"]),
int(row["Lower right corner X"]),
int(row["Lower right corner Y"]),
category=".".join([supercategory, row["Annotation tag"]]),
)
data.label.box2d.append(label)
def _load_labels(label_file: str) -> List[LabeledBox2D]:
with open(label_file, "r", encoding="utf-8") as fp:
objects = xmltodict.parse(fp.read())["annotation"]["object"]
box2ds = []
if not isinstance(objects, list):
objects = [objects]
for obj in objects:
bndbox = obj["bndbox"]
box2ds.append(
LabeledBox2D(
xmin=float(bndbox["xmin"]),
ymin=float(bndbox["ymin"]),
xmax=float(bndbox["xmax"]),
ymax=float(bndbox["ymax"]),
category=obj["name"],
))
return box2ds
def _load_positive_segment(segment_name: str, segment_path: str) -> Segment:
if segment_name.startswith("vid"):
# Pad zero for segment name to change "vid0" to "vid00"
segment_name = f"{segment_name[:3]}{int(segment_name[3:]):02}"
segment = Segment(segment_name)
annotation_file = glob(
os.path.join(segment_path, "frameAnnotations-*",
"frameAnnotations.csv"))[0]
image_folder = os.path.dirname(annotation_file)
pre_filename = ""
with open(annotation_file, "r", encoding="utf-8") as fp:
for annotation in csv.DictReader(fp, delimiter=";"):
filename = annotation["Filename"]
if filename != pre_filename:
data = Data(os.path.join(image_folder, filename))
data.label.box2d = []
segment.append(data)
pre_filename = filename
occluded, on_another_road = annotation[
"Occluded,On another road"].split(",", 1)
data.label.box2d.append(
LabeledBox2D(
int(annotation["Upper left corner X"]),
int(annotation["Upper left corner Y"]),
int(annotation["Lower right corner X"]),
int(annotation["Lower right corner Y"]),
category=annotation["Annotation tag"],
attributes={
"Occluded":
bool(int(occluded)),
"On another road":
bool(int(on_another_road)),
"Origin file":
annotation["Origin file"],
"Origin frame number":
int(annotation["Origin frame number"]),
"Origin track":
annotation["Origin track"],
"Origin track frame number":
int(annotation["Origin track frame number"]),
},
))
return segment
def _get_data_part2(root_path: str, aniamls: Iterable[str]) -> Iterator[Data]:
try:
import xmltodict # pylint: disable=import-outside-toplevel
except ModuleNotFoundError as error:
raise ModuleImportError(module_name=error.name) from error
for animal in aniamls:
for image_path in glob(
os.path.join(root_path, "animalpose_image_part2", animal,
"*.jpeg")):
data = Data(
image_path,
target_remote_path=f"{animal}/{os.path.basename(image_path)}")
annotation_path = os.path.join(
root_path,
"animalpose_anno2",
animal,
f"{os.path.splitext(os.path.basename(image_path))[0]}.xml",
)
with open(annotation_path, encoding="utf-8") as fp:
labels: Any = xmltodict.parse(fp.read())
box2d = labels["annotation"]["visible_bounds"]
data.label.box2d = [
LabeledBox2D.from_xywh(
x=float(box2d["@xmin"]),
y=float(
box2d["@xmax"]), # xmax means ymin in the annotation
width=float(box2d["@width"]),
height=float(box2d["@height"]),
category=animal,
)
]
keypoints2d = LabeledKeypoints2D(category=animal)
for keypoint in labels["annotation"]["keypoints"]["keypoint"]:
keypoints2d.append(
Keypoint2D(float(keypoint["@x"]), float(keypoint["@y"]),
int(keypoint["@visible"])))
data.label.keypoints2d = [keypoints2d]
yield data
def HeadPoseImage(path: str) -> Dataset:
"""`Head Pose Image <http://crowley-coutaz.fr\
/Head%20Pose%20Image%20Database.html>`_ dataset.
The file structure should be like::
<path>
Person01/
person01100-90+0.jpg
person01100-90+0.txt
person01101-60-90.jpg
person01101-60-90.txt
...
Person02/
Person03/
...
Person15/
Arguments:
path: The root directory of the dataset.
Returns:
Loaded :class:`~tensorbay.dataset.dataset.Dataset` instance.
"""
dataset = Dataset(DATASET_NAME)
dataset.load_catalog(
os.path.join(os.path.dirname(__file__), "catalog.json"))
segment = dataset.create_segment()
image_paths = glob(os.path.join(path, "Person*", "*.jpg"))
for image_path in image_paths:
image_name = os.path.basename(image_path)
data = Data(image_path)
data.label.box2d = [
LabeledBox2D(
*_load_label_box(image_path.replace("jpg", "txt")),
category=image_name[6:8],
attributes=_load_attributes(image_name),
)
]
segment.append(data)
return dataset
def test_init(self):
labeledbox2d = LabeledBox2D(
1,
2,
4,
6,
category="cat",
attributes={"gender": "male"},
instance="12345",
)
assert labeledbox2d.category == "cat"
assert labeledbox2d.attributes == {"gender": "male"}
assert labeledbox2d.instance == "12345"
assert labeledbox2d[0] == 1
assert labeledbox2d[1] == 2
assert labeledbox2d[2] == 4
assert labeledbox2d[3] == 6
def _get_mot_data(image_path: str, label_content: Dict[str, Any]) -> Data:
data = Data(image_path)
labeled_box2ds = []
for label_info in label_content.get("labels", ()):
box2d_info = label_info.get("box2d")
if not box2d_info:
continue
labeled_box2d = LabeledBox2D(
box2d_info["x1"],
box2d_info["y1"],
box2d_info["x2"],
box2d_info["y2"],
category=label_info["category"],
attributes=label_info["attributes"],
instance=label_info["id"],
)
labeled_box2ds.append(labeled_box2d)
data.label.box2d = labeled_box2ds
return data
def test_from_xywh(self):
x, y, width, height = 1, 2, 3, 4
xmin, xmax, ymin, ymax = 1, 2, 4, 6
labeledbox2d = LabeledBox2D.from_xywh(
x,
y,
width,
height,
category="cat",
attributes={"gender": "male"},
instance="12345",
)
assert labeledbox2d.category == "cat"
assert labeledbox2d.attributes == {"gender": "male"}
assert labeledbox2d.instance == "12345"
assert labeledbox2d[0] == xmin
assert labeledbox2d[1] == xmax
assert labeledbox2d[2] == ymin
assert labeledbox2d[3] == ymax
def get_voc_detection_data(
stem: str, image_path: str, annotation_path: str, boolean_attributes: List[str]
) -> Data:
"""Get all information of the datum corresponding to voc-like label files.
Arguments:
stem: The filename without extension of the data.
image_path: The path of the image directory.
annotation_path: The path of the annotation directory.
boolean_attributes: The list of boolean attribute.
Returns:
Data: class:`~tensorbay.dataset.data.Data` instance.
"""
data = Data(os.path.join(image_path, f"{stem}.jpg"))
box2d = []
with open(os.path.join(annotation_path, f"{stem}.xml"), encoding="utf-8") as fp:
labels: Any = xmltodict.parse(fp.read())
objects = labels["annotation"]["object"]
if not isinstance(objects, list):
objects = [objects]
for obj in objects:
attributes = {attribute: bool(int(obj[attribute])) for attribute in boolean_attributes}
attributes["pose"] = obj["pose"]
bndbox = obj["bndbox"]
box2d.append(
LabeledBox2D(
float(bndbox["xmin"]),
float(bndbox["ymin"]),
float(bndbox["xmax"]),
float(bndbox["ymax"]),
category=obj["name"],
attributes=attributes,
)
)
data.label.box2d = box2d
return data
def _get_data(path: str, annotations: Any,
flag: bool) -> Iterator[Tuple[Data, str]]:
filepath_to_data: Dict[str, Data] = {}
for annotation in annotations:
filepath = annotation["filepath"][0]
keypoints = LabeledKeypoints2D(
annotation["coords"].T[_VALID_KEYPOINT_INDICES],
attributes={
"poselet_hit_idx": annotation["poselet_hit_idx"].T.tolist()
},
)
box2d = LabeledBox2D(*annotation["torsobox"][0].tolist())
if filepath not in filepath_to_data:
data = Data(os.path.join(path, "images", filepath))
data.label.keypoints2d = [keypoints]
data.label.box2d = [box2d]
attribute = {"currframe": int(annotation["currframe"][0][0])}
if flag:
attribute["isunchecked"] = bool(annotation["isunchecked"])
data.label.classification = Classification(
category=annotation["moviename"][0], attributes=attribute)
filepath_to_data[filepath] = data
if annotation["istrain"]:
segment_name = "train"
elif annotation["istest"]:
segment_name = "test"
else:
segment_name = "bad"
yield data, segment_name
else:
image_data = filepath_to_data[filepath]
image_data.label.keypoints2d.append(keypoints)
image_data.label.box2d.append(box2d)
def _generate_data(image_path: str, labels: Dict[str, Any]) -> Data:
data = Data(image_path)
data.label.box2d = []
image_id = labels["image_name_id_map"][os.path.basename(image_path)]
image_annotations_map = labels["image_annotations_map"]
if image_id not in image_annotations_map:
return data
annotations = labels["annotations"]
poses = labels["poses"]
categories = labels["categories"]
for annotation_id in image_annotations_map[image_id]:
annotation = annotations[annotation_id]
x_top, y_top, width, height = annotation["bbox"]
attributes = {
"occluded": annotation["occluded"],
"difficult": annotation["difficult"],
"pose": poses[annotation["pose_id"] - 1]["name"],
"truncated": annotation["truncated"],
}
data.label.box2d.append(
LabeledBox2D.from_xywh(
x=x_top,
y=y_top,
width=width,
height=height,
category=categories[annotation["category_id"]]["name"],
attributes=attributes,
instance=str(annotation["tracking_id"]),
))
return data
def _get_box2d_and_panoptic_mask(
objects: Any,
original_mask_path: str,
new_mask_path: str,
category_ids: Dict[str, int],
) -> Tuple[List[LabeledBox2D], PanopticMask]:
all_category_ids: Dict[int, int] = {}
original_mask = np.array(Image.open(original_mask_path))
box2ds: List[LabeledBox2D] = []
if not isinstance(objects, list):
objects = [objects]
rgba_to_instance_id: Dict[Tuple[int, ...], int] = {}
for index, obj in enumerate(objects, 1):
category = ".".join(
takewhile(bool, (obj.get(f"category{i}", "")
for i in range(_MAX_CATEGORY_LEVEL))))
bndbox = obj["bndbox2D"]
box2ds.append(
LabeledBox2D(
int(bndbox["xmin"]),
int(bndbox["ymin"]),
int(bndbox["xmax"]),
int(bndbox["ymax"]),
category=category,
attributes={"focus_blur": bndbox["focus_blur"]},
))
rgba_to_instance_id[tuple(
map(int, obj["object_mask_color_rgba"].split(",")))] = index
all_category_ids[index] = category_ids[category]
mask = np.vectorize(lambda r, g, b: rgba_to_instance_id.get(
(r, g, b, 255), 0))(original_mask[:, :, 0], original_mask[:, :, 1],
original_mask[:, :, 2]).astype(np.uint8)
Image.fromarray(mask).save(new_mask_path)
panoptic_mask = PanopticMask(new_mask_path)
panoptic_mask.all_category_ids = all_category_ids
return box2ds, panoptic_mask
def test_loads(self):
contents = {
"box2d": {
"xmin": 1,
"ymin": 2,
"xmax": 5,
"ymax": 8
},
"category": "cat",
"attributes": {
"gender": "male"
},
"instance": 12345,
}
labeledbox2d = LabeledBox2D.loads(contents)
assert labeledbox2d.category == "cat"
assert labeledbox2d.attributes == {"gender": "male"}
assert labeledbox2d.instance == 12345
assert labeledbox2d[0] == 1
assert labeledbox2d[1] == 2
assert labeledbox2d[2] == 5
assert labeledbox2d[3] == 8
def _get_box2ds(label_path: str) -> Dict[str, LabeledBox2D]:
all_box2d_attributes: DefaultDict[str, AttributeType] = defaultdict(dict)
# The normal format of each line of the file is
# n000002/0032_01.jpg 0
# n000002/0039_01.jpg 0
# n000002/0090_01.jpg 0
# ...
for file_path in glob(os.path.join(label_path, "attributes", "*.txt")):
attribute_name = os.path.basename(file_path).rstrip(".txt").split(
"-", 1)[1]
with open(file_path, encoding="utf-8") as fp:
for line in fp:
name, attribute_value = line.rstrip("\n").split("\t", 1)
all_box2d_attributes[name.rstrip(
".jpg")][attribute_name] = bool(int(attribute_value))
all_boxes = {}
for file_path in (
os.path.join(label_path, "loose_bb_test.csv"),
os.path.join(label_path, "loose_bb_train.csv"),
):
# The normal format of each line of the file is
# NAME_ID,X,Y,W,H
# "n000001/0001_01",60,60,79,109
# "n000001/0002_01",134,81,207,295
# "n000001/0003_01",58,32,75,103
# ...
with open(file_path, encoding="utf-8") as fp:
for row in islice(csv.reader(fp), 1, None):
name_id = row.pop(0).strip('"')
box = LabeledBox2D.from_xywh(*map(float, row))
box2d_attribute = all_box2d_attributes.get(name_id)
if box2d_attribute:
box.attributes = box2d_attribute
all_boxes[name_id] = box
return all_boxes
