download_data.py
Python script, ASCII text executable
1
import os
2
import shutil
3
import httpx
4
import tqdm
5
from PIL import Image, ImageOps
6
import tempfile
7
from pathlib import Path
8
9
query_yaml = """
10
want:
11
- has: ["waste"]
12
- nature: ["photo"]
13
"""
14
15
object_types = httpx.get("https://datasets.roundabout-host.com/api/object-types").json()
16
17
base_types = {
18
"plastic household waste": set(),
19
"glass household waste": set(),
20
"metal household waste": set(),
21
"paper and cardboard": set(),
22
"organic household waste": set(),
23
"household waste": set(),
24
"waste": set(),
25
}
26
27
def get_objects(object_info):
28
objects = set()
29
for object in object_info["children"]:
30
#print(f"Scanning {object}")
31
info = httpx.get("https://datasets.roundabout-host.com/api/object/" + object).json()
32
objects |= get_objects(info)
33
objects.add(object_info["id"])
34
return objects
35
36
for object_name in base_types:
37
object_info = httpx.get("https://datasets.roundabout-host.com/api/object/" + object_name).json()
38
base_types[object_name] = get_objects(object_info)
39
print(object_name, base_types[object_name])
40
41
result = {"resources": True} # dummy value to enter the loop
42
photos = []
43
offset = 0
44
limit = 192
45
output = Path("data/")
46
47
print("Downloading photo metadata...")
48
49
while result["resources"]:
50
result = httpx.post(f"https://datasets.roundabout-host.com/api/query-pictures?offset={offset}&limit={limit}", data={"query": query_yaml}).json()
51
photos += result["resources"]
52
offset += limit
53
print(f"Received photos {offset-limit} to {offset}")
54
55
shutil.rmtree(output, ignore_errors=True)
56
os.makedirs(output)
57
58
class_mapping = {
59
"plastic household waste": 0,
60
"glass household waste": 1,
61
"metal household waste": 2,
62
"paper and cardboard": 3,
63
"organic household waste": 4,
64
"household waste": 5,
65
"waste": 6,
66
}
67
68
with tempfile.NamedTemporaryFile(delete_on_close=False, mode="wb") as temporary_file:
69
for photo in tqdm.tqdm(photos):
70
# Download the photo
71
result = httpx.get(photo["download"], follow_redirects=True)
72
temporary_file.write(result.content)
73
temporary_file.seek(0)
74
image = Image.open(temporary_file.name, formats=["JPEG"])
75
ImageOps.exif_transpose(image, in_place=True)
76
image.thumbnail((640, 640))
77
image.save(output / (str(photo["id"]) + ".jpg"))
78
# Download the annotations
79
with open(output / (str(photo["id"]) + ".txt"), "w") as annotation_file:
80
for region in photo["regions"]:
81
klass = -1
82
for base_type in base_types:
83
if region["object"] in base_types[base_type]:
84
klass = class_mapping[base_type]
85
break
86
if klass == -1:
87
continue
88
89
if region["type"] == "bbox":
90
cx = region["shape"]["x"] + region["shape"]["w"] / 2
91
cy = region["shape"]["y"] + region["shape"]["h"] / 2
92
w = region["shape"]["w"]
93
h = region["shape"]["h"]
94
elif region["type"] == "polygon":
95
x = [point["x"] for point in region["shape"]]
96
y = [point["y"] for point in region["shape"]]
97
cx = (min(x) + max(x)) / 2
98
cy = (min(y) + max(y)) / 2
99
w = max(x) - min(x)
100
h = max(y) - min(y)
101
102
annotation_file.write(f"{klass} {cx} {cy} {w} {h}\n")
103
104