#!/usr/bin/env python3
#
# Copyright 2021 Graviti. Licensed under MIT License.
#
# pylint: disable=invalid-name
"""Dataloader of JHU-CROWD dataset."""
import os
from typing import Dict, List
from tensorbay.dataset import Data, Dataset
from tensorbay.label import Classification, LabeledBox2D
from tensorbay.opendataset._utility import glob
SEGMENT_LIST = ["train", "val", "test"]
DATASET_NAME = "JHU-CROWD"
_OCCLUSION_MAP = {1: "visible", 2: "partial-occlusion", 3: "full-occlusion"}
_WEATHER_CONDITION_MAP = {0: "no weather degradationi", 1: "fog/haze", 2: "rain", 3: "snow"}
[docs]def JHU_CROWD(path: str) -> Dataset:
"""`JHU-CROWD++ <http://www.crowd-counting.com/>`_ dataset.
The file structure should be like::
<path>
train/
images/
0000.jpg
...
gt/
0000.txt
...
image_labels.txt
test/
val/
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"))
for segment_name in SEGMENT_LIST:
segment = dataset.create_segment(segment_name)
segment_path = os.path.join(path, segment_name)
image_root_path = os.path.join(segment_path, "images")
image_paths = glob(os.path.join(image_root_path, "*.jpg"))
image_labels = _load_image_labels(os.path.join(segment_path, "image_labels.txt"))
for image_path in image_paths:
data = Data(image_path)
image_file = os.path.basename(image_path)
label_file = image_file.replace("jpg", "txt")
data.label.box2d = _load_box_labels(os.path.join(segment_path, "gt", label_file))
data.label.classification = image_labels[os.path.splitext(image_file)[0]]
segment.append(data)
return dataset
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 _load_image_labels(file_path: str) -> Dict[str, Classification]:
with open(file_path, encoding="utf-8") as fp:
image_labels = {}
for line in fp:
img_index, count, scene, weather, distractor = line.strip().split(",")
attributes = {
"total-count": int(count),
"scene-type": scene,
"weather-condition": _WEATHER_CONDITION_MAP[int(weather)],
"distractor": bool(int(distractor)),
}
image_labels[img_index] = Classification(attributes=attributes)
return image_labels