feat(opendataset): add dataloader for UAVDT dataset

PR Closed: #723

tensorbay/opendataset/UAVDT/__init__.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+#
+# Copytright 2020 Graviti. Licensed under MIT License.
+#
+# pylint: disable=invalid-name
+
+"""Dataloader of UAVDT."""
+
+from .loader import UAVDT
+
+__all__ = ["UAVDT"]

tensorbay/opendataset/UAVDT/catalog.json

@@ -0,0 +1,62 @@
+{
+    "BOX2D": {
+        "isTracking": true,
+        "categories": [{ "name": "car" }, { "name": "trunck" }, { "name": "bus" }],
+        "attributes": [
+            {
+                "name": "occlusion",
+                "description": "The score in the GROUNDTRUTH file indicates the degree of object parts appears outside a frame",
+                "enum": ["no", "large", "medium", "small"]
+            },
+            {
+                "name": "out_of_view",
+                "description": "The score in the GROUNDTRUTH fileindicates the fraction of objects being occluded",
+                "enum": ["no", "medium", "small"]
+            }
+        ]
+    },
+    "CLASSIFICATION": {
+        "attributes": [
+            {
+                "name": "daylight",
+                "type": "boolean"
+            },
+            {
+                "name": "night",
+                "type": "boolean"
+            },
+            {
+                "name": "fog",
+                "type": "boolean"
+            },
+            {
+                "name": "low-alt",
+                "type": "boolean"
+            },
+            {
+                "name": "medium-alt",
+                "type": "boolean"
+            },
+            {
+                "name": "high-alt",
+                "type": "boolean"
+            },
+            {
+                "name": "front-view",
+                "type": "boolean"
+            },
+            {
+                "name": "side-view",
+                "type": "boolean"
+            },
+            {
+                "name": "bird-view",
+                "type": "boolean"
+            },
+            {
+                "name": "long-term",
+                "type": "boolean"
+            }
+        ]
+    }
+}

tensorbay/opendataset/UAVDT/loader.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+#
+# Copytright 2020 Graviti. Licensed under MIT License.
+#
+# pylint: disable=invalid-name, missing-module-docstring
+
+import csv
+import os
+from collections import defaultdict
+from typing import Dict, List
+
+from ...dataset import Data, Dataset
+from ...label import Box2DSubcatalog, Classification, ClassificationSubcatalog, LabeledBox2D
+from .._utility import glob
+
+_DATASET_NAME = "UAVDT"
+
+_MODES = ("train", "test")
+
+_FIELDNAMES = (
+    "frame_index",
+    "target_id",
+    "bbox_left",
+    "bbox_top",
+    "bbox_width",
+    "bbox_height",
+    "out_of_view",
+    "occlusion",
+    "object_category",
+)
+
+_XYWH_KEYS = _FIELDNAMES[2:6]
+
+
+def UAVDT(path: str) -> Dataset:
+    """Dataloader of the `UAVDT`_ Dataset.
+
+    The "score", "in-view", "occlusion" fields in MOT Groundtruth file(*_gt.txt) are constant, and
+    other fields in that file are the same with such fields in DET Groundtruth file(*_gt_whole.txt).
+    Therefore, they are not included in the dataloader.
+
+    The Ignore Areas file(*_gt_ignore.txt) is useless, so they are not included in the dataloader
+    neither.
+
+    .. _UAVDT: https://sites.google.com/site/daviddo0323/projects/uavdt
+
+    The file structure of UAVDT looks like:
+
+        <path>
+            M_attr/
+                test/
+                    M0203_attr.txt
+                    ...
+                train/
+                    M0101_attr.txt
+                    ...
+            UAVDT_Benchmark_M/
+                M0101/
+                    img000001.jpg
+                    ...
+                ...
+            UAV-benchmark-MOTD_v1.0/
+                GT/
+                    M0101_gt_ignore.txt
+                    M0101_gt.txt
+                    M0101_gt_whole.txt
+                    ...
+
+    Arguments:
+        path: The root directory of the dataset.
+
+    Returns:
+        Loaded :class:`~tensorbay.dataset.dataset.Dataset` instance.
+
+    """
+    root_path = os.path.abspath(os.path.expanduser(path))
+
+    dataset = Dataset(_DATASET_NAME)
+    dataset.notes.is_continuous = True
+    dataset.load_catalog(os.path.join(os.path.dirname(__file__), "catalog.json"))
+
+    for mode in _MODES:
+        for sequence_attributes_path in glob(os.path.join(root_path, "M_attr", mode, "*.txt")):
+            sequence = os.path.basename(sequence_attributes_path)[:5]
+            segment = dataset.create_segment(f"{mode}-{sequence}")
+
+            classification = _extract_classification(
+                sequence_attributes_path, dataset.catalog.classification
+            )
+
+            frame_id_box2d_map = _extract_box2d(root_path, sequence, dataset.catalog.box2d)
+
+            image_paths = glob(os.path.join(root_path, "UAV-benchmark-M", sequence, "*.jpg"))
+            for image_path in image_paths:
+                data = Data(image_path)
+                data.label.classification = classification
+                # The image_name looks like `img<frame_id>.jpg`.
+                # The frame_id is consist of six digital numbers.
+                data.label.box2d = frame_id_box2d_map[int(os.path.basename(image_path)[3:9])]
+                segment.append(data)
+
+    return dataset
+
+
+def _extract_classification(
+    path: str, classification_catalog: ClassificationSubcatalog
+) -> Classification:
+    with open(path) as fp:
+        attribute_names = classification_catalog.attributes.keys()
+        csv_reader = csv.reader(fp)
+        elements = next(csv_reader)
+        attributes = {
+            attribute_name: bool(int(value))
+            for attribute_name, value in zip(attribute_names, elements)
+        }
+
+    return Classification(attributes=attributes)
+
+
+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")) 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

tensorbay/opendataset/__init__.py

@@ -40,6 +40,7 @@
 from .THCHS30 import THCHS30
 from .THUCNews import THUCNews
 from .TLR import TLR
+from .UAVDT import UAVDT
 from .WIDER_FACE import WIDER_FACE
 
 __all__ = [
@@ -82,5 +83,6 @@
     "THCHS30",
     "THUCNews",
     "TLR",
+    "UAVDT",
     "WIDER_FACE",
 ]