-
Notifications
You must be signed in to change notification settings - Fork 2
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.
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
149 implement robust dpatch attack #155
Changes from all commits
a8b7ec2
5785160
19c5a55
6b52a4f
5f21e80
d81db21
09e6a8e
9d29436
835b906
78e1182
9b6f67e
b68eca5
1ba82ca
bc68b4f
3aef3d0
538b634
eb0eff2
d124aa2
5121d35
ba2ea7c
ff23c18
0a0d9a8
47110ed
e9ce9bd
373d2eb
10ce3c1
0056c67
fb32dfe
dd59429
509cb8b
6dc7043
7e80202
062e285
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,226 @@ | ||
"""Utilities to load the VisDrone 2019 dataset.""" | ||
|
||
import csv | ||
import io | ||
from pathlib import Path | ||
from pprint import pprint | ||
from typing import Any, Dict, Iterator, List, Tuple | ||
|
||
import albumentations as A | ||
import albumentations.pytorch | ||
import datasets | ||
import numpy as np | ||
|
||
import armory.data | ||
import armory.dataset | ||
|
||
|
||
def create_dataloader( | ||
dataset: datasets.Dataset, max_size: int, **kwargs | ||
) -> armory.dataset.ObjectDetectionDataLoader: | ||
""" | ||
Create an Armory object detection dataloader for the given VisDrone2019 dataset split. | ||
|
||
Args: | ||
dataset: VisDrone2019 dataset split | ||
max_size: Maximum image size to which to resize and pad image samples | ||
**kwargs: Additional keyword arguments to pass to the dataloader constructor | ||
|
||
Return: | ||
Armory object detection dataloader | ||
""" | ||
resize = A.Compose( | ||
[ | ||
A.LongestMaxSize(max_size=max_size), | ||
A.PadIfNeeded( | ||
min_height=max_size, | ||
min_width=max_size, | ||
border_mode=0, | ||
value=(0, 0, 0), | ||
), | ||
A.ToFloat(max_value=255), | ||
albumentations.pytorch.ToTensorV2(), | ||
], | ||
bbox_params=A.BboxParams( | ||
format="coco", | ||
label_fields=["id", "category", "occlusion", "truncation"], | ||
), | ||
) | ||
|
||
def transform(sample): | ||
tmp = dict(**sample) | ||
treubig26 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
tmp["image"] = [] | ||
tmp["objects"] = [] | ||
for image, objects in zip(sample["image"], sample["objects"]): | ||
res = resize( | ||
image=np.asarray(image), | ||
bboxes=objects["bbox"], | ||
id=objects["id"], | ||
category=objects["category"], | ||
occlusion=objects["occlusion"], | ||
truncation=objects["truncation"], | ||
) | ||
tmp["image"].append(res.pop("image")) | ||
tmp["objects"].append(res) | ||
return tmp | ||
|
||
dataset.set_transform(transform) | ||
|
||
return armory.dataset.ObjectDetectionDataLoader( | ||
dataset, | ||
image_key="image", | ||
dim=armory.data.ImageDimensions.CHW, | ||
scale=armory.data.Scale( | ||
dtype=armory.data.DataType.FLOAT, | ||
max=1.0, | ||
), | ||
objects_key="objects", | ||
boxes_key="bboxes", | ||
format=armory.data.BBoxFormat.XYWH, | ||
labels_key="category", | ||
**kwargs, | ||
) | ||
|
||
|
||
TRAIN_URL = "https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-train.zip" | ||
VAL_URL = "https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-val.zip" | ||
TEST_URL = "https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-dev.zip" | ||
|
||
treubig26 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
def load_dataset() -> datasets.DatasetDict: | ||
""" | ||
Load the train and validation splits of the VisDrone2019 dataset. | ||
|
||
Return: | ||
Dictionary containing the train and validation splits | ||
""" | ||
dl_manager = datasets.DownloadManager(dataset_name="VisDrone2019") | ||
ds_features = features() | ||
paths = dl_manager.download({"train": TRAIN_URL, "val": VAL_URL, "test": TEST_URL}) | ||
train_files = dl_manager.iter_archive(paths["train"]) | ||
val_files = dl_manager.iter_archive(paths["val"]) | ||
test_files = dl_manager.iter_archive(paths["test"]) | ||
return datasets.DatasetDict( | ||
{ | ||
datasets.Split.TRAIN: datasets.Dataset.from_generator( | ||
generate_samples, | ||
gen_kwargs={"files": train_files}, | ||
features=ds_features, | ||
), | ||
datasets.Split.VALIDATION: datasets.Dataset.from_generator( | ||
generate_samples, | ||
gen_kwargs={"files": val_files}, | ||
features=ds_features, | ||
), | ||
treubig26 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
datasets.Split.TEST: datasets.Dataset.from_generator( | ||
generate_samples, | ||
gen_kwargs={"files": test_files}, | ||
features=ds_features, | ||
), | ||
} | ||
) | ||
|
||
|
||
CATEGORIES = [ | ||
# The YOLOv5 model removed this class and shifted all others down by 1 when | ||
# it trained on the VisDrone data | ||
# "ignored", | ||
Comment on lines
+125
to
+127
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Let's make this dataset useable for other trained models that might keep original category labels. |
||
"pedestrian", | ||
"people", | ||
"bicycle", | ||
"car", | ||
"van", | ||
"truck", | ||
"tricycle", | ||
"awning-tricycle", | ||
"bus", | ||
"motor", | ||
# The YOLOv5 model also ignored/removed this class | ||
# "other", | ||
] | ||
|
||
|
||
def features() -> datasets.Features: | ||
"""Create VisDrone2019 dataset features""" | ||
return datasets.Features( | ||
{ | ||
"image_id": datasets.Value("int64"), | ||
"file_name": datasets.Value("string"), | ||
"image": datasets.Image(), | ||
"objects": datasets.Sequence( | ||
{ | ||
"id": datasets.Value("int64"), | ||
"bbox": datasets.Sequence(datasets.Value("float32"), length=4), | ||
"category": datasets.ClassLabel( | ||
num_classes=len(CATEGORIES), names=CATEGORIES | ||
), | ||
"truncation": datasets.Value("int32"), | ||
"occlusion": datasets.Value("int32"), | ||
} | ||
), | ||
} | ||
) | ||
|
||
|
||
ANNOTATION_FIELDS = [ | ||
"x", | ||
"y", | ||
"width", | ||
"height", | ||
"score", | ||
"category_id", | ||
"truncation", | ||
"occlusion", | ||
] | ||
|
||
|
||
def load_annotations(file: io.BufferedReader) -> List[Dict[str, Any]]: | ||
"""Load annotations/objects from the given file""" | ||
reader = csv.DictReader( | ||
io.StringIO(file.read().decode("utf-8")), fieldnames=ANNOTATION_FIELDS | ||
) | ||
annotations = [] | ||
for idx, row in enumerate(reader): | ||
score = int(row["score"]) | ||
category = int(row["category_id"]) | ||
if score != 0: # Drop annotations with score of 0 (class-0 & class-11) | ||
category -= 1 # The model was trained with 0-indexed categories starting at pedestrian | ||
Comment on lines
+186
to
+187
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Remapping the category labels should be a parameter of |
||
bbox = list(map(float, [row[k] for k in ANNOTATION_FIELDS[:4]])) | ||
if bbox[2] == 0 or bbox[3] == 0: | ||
continue | ||
annotations.append( | ||
{ | ||
"id": idx, | ||
"bbox": bbox, | ||
"category": category, | ||
"truncation": row["truncation"], | ||
"occlusion": row["occlusion"], | ||
} | ||
) | ||
return annotations | ||
|
||
|
||
def generate_samples( | ||
files: Iterator[Tuple[str, io.BufferedReader]], annotation_file_ext: str = ".txt" | ||
) -> Iterator[Dict[str, Any]]: | ||
"""Generate dataset samples from the given files in a VisDrone2019 archive""" | ||
annotations = {} | ||
images = {} | ||
for path, file in files: | ||
file_name = Path(path).stem | ||
if Path(path).suffix == annotation_file_ext: | ||
annotations[file_name] = load_annotations(file) | ||
else: | ||
images[file_name] = {"path": path, "bytes": file.read()} | ||
|
||
for idx, (file_name, annotation) in enumerate(annotations.items()): | ||
yield { | ||
"image_id": idx, | ||
"file_name": file_name, | ||
"image": images[file_name], | ||
"objects": annotation, | ||
} | ||
|
||
|
||
if __name__ == "__main__": | ||
pprint(load_dataset()) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It may not be necessary to pad all input images to the max size, some models take a list of images as input (e.g. Faster-RCNN).