Skip to content

Commit 9f7929e

Browse files
WovchenaWovchena
authored
Async tests/python/accuracy/prepare_data.py (#91)
1 parent 16a2a94 commit 9f7929e

File tree

2 files changed

+101
-66
lines changed

2 files changed

+101
-66
lines changed

‎model_api/python/setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@
3838
install_requires=(SETUP_DIR / "requirements.txt").read_text(),
3939
extras_require={
4040
"ovms": (SETUP_DIR / "requirements_ovms.txt").read_text(),
41-
"tests": ["pytest", "openvino-dev[onnx,pytorch,tensorflow2]"],
41+
"tests": ["httpx", "pytest", "openvino-dev[onnx,pytorch,tensorflow2]"],
4242
},
4343
long_description=(SETUP_DIR.parents[1] / "README.md").read_text(),
4444
long_description_content_type="text/markdown",

‎tests/python/accuracy/prepare_data.py

Lines changed: 100 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -1,81 +1,116 @@
11
import argparse
2-
import os
3-
from urllib.request import urlopen, urlretrieve
2+
import asyncio
3+
from io import BytesIO
4+
from pathlib import Path
5+
from zipfile import ZipFile
46

7+
import httpx
58

6-
def retrieve_otx_model(data_dir, model_name, format="xml"):
7-
destination_folder = os.path.join(data_dir, "otx_models")
8-
os.makedirs(destination_folder, exist_ok=True)
9-
if format == "onnx":
10-
urlretrieve(
11-
f"https://storage.openvinotoolkit.org/repositories/model_api/test/otx_models/{model_name}/model.onnx",
12-
f"{destination_folder}/{model_name}.onnx",
13-
)
14-
else:
15-
urlretrieve(
16-
f"https://storage.openvinotoolkit.org/repositories/model_api/test/otx_models/{model_name}/openvino.xml",
17-
f"{destination_folder}/{model_name}.xml",
18-
)
19-
urlretrieve(
20-
f"https://storage.openvinotoolkit.org/repositories/model_api/test/otx_models/{model_name}/openvino.bin",
21-
f"{destination_folder}/{model_name}.bin",
22-
)
239

10+
async def download_images(data_dir):
11+
async with httpx.AsyncClient(timeout=20.0) as client:
12+
COCO128_URL = "https://ultralytics.com/assets/coco128.zip"
13+
archive = await client.get(COCO128_URL, follow_redirects=True)
14+
with ZipFile(BytesIO(archive.content)) as zfile:
15+
zfile.extractall(data_dir)
16+
image = await client.get(
17+
"https://raw.githubusercontent.com/Shenggan/BCCD_Dataset/master/BCCD/JPEGImages/BloodImage_00007.jpg"
18+
)
19+
with data_dir / "BloodImage_00007.jpg" as im:
20+
im.write_bytes(image.content)
2421

25-
def prepare_data(data_dir="./data"):
26-
from io import BytesIO
27-
from zipfile import ZipFile
2822

29-
COCO128_URL = "https://ultralytics.com/assets/coco128.zip"
23+
async def stream_file(client, url, filename):
24+
async with client.stream("GET", url) as stream:
25+
with open(filename, "wb") as file:
26+
async for data in stream.aiter_bytes():
27+
file.write(data)
3028

31-
with urlopen(COCO128_URL) as zipresp:
32-
with ZipFile(BytesIO(zipresp.read())) as zfile:
33-
zfile.extractall(data_dir)
3429

35-
urlretrieve(
36-
"https://raw.githubusercontent.com/Shenggan/BCCD_Dataset/master/BCCD/JPEGImages/BloodImage_00007.jpg",
37-
os.path.join(data_dir, "BloodImage_00007.jpg"),
38-
)
30+
async def download_otx_model(client, otx_models_dir, model_name, format="xml"):
31+
if format == "onnx":
32+
await stream_file(
33+
client,
34+
f"https://storage.openvinotoolkit.org/repositories/model_api/test/otx_models/{model_name}/model.onnx",
35+
f"{otx_models_dir}/{model_name}.onnx",
36+
)
37+
else:
38+
await asyncio.gather(
39+
stream_file(
40+
client,
41+
f"https://storage.openvinotoolkit.org/repositories/model_api/test/otx_models/{model_name}/openvino.xml",
42+
f"{otx_models_dir}/{model_name}.xml",
43+
),
44+
stream_file(
45+
client,
46+
f"https://storage.openvinotoolkit.org/repositories/model_api/test/otx_models/{model_name}/openvino.bin",
47+
f"{otx_models_dir}/{model_name}.bin",
48+
),
49+
)
3950

4051

41-
if __name__ == "__main__":
42-
parser = argparse.ArgumentParser(description="Data and model prepare script")
52+
async def main():
53+
parser = argparse.ArgumentParser()
4354
parser.add_argument(
4455
"-d",
45-
dest="data_dir",
46-
default="./data",
56+
"--data_dir",
57+
type=Path,
58+
required=True,
4759
help="Directory to store downloaded models and datasets",
4860
)
49-
5061
args = parser.parse_args()
5162

52-
prepare_data(args.data_dir)
53-
retrieve_otx_model(args.data_dir, "mlc_mobilenetv3_large_voc")
54-
retrieve_otx_model(args.data_dir, "mlc_efficient_b0_voc")
55-
retrieve_otx_model(args.data_dir, "mlc_efficient_v2s_voc")
56-
retrieve_otx_model(args.data_dir, "det_mobilenetv2_atss_bccd")
57-
retrieve_otx_model(args.data_dir, "det_mobilenetv2_atss_bccd_onnx", "onnx")
58-
retrieve_otx_model(args.data_dir, "cls_mobilenetv3_large_cars")
59-
retrieve_otx_model(args.data_dir, "cls_mobilenetv3_large_cars", "onnx")
60-
retrieve_otx_model(args.data_dir, "cls_efficient_b0_cars")
61-
retrieve_otx_model(args.data_dir, "cls_efficient_v2s_cars")
62-
retrieve_otx_model(args.data_dir, "Lite-hrnet-18")
63-
retrieve_otx_model(args.data_dir, "Lite-hrnet-18_mod2")
64-
retrieve_otx_model(args.data_dir, "Lite-hrnet-s_mod2")
65-
retrieve_otx_model(args.data_dir, "Lite-hrnet-s_mod2", "onnx")
66-
retrieve_otx_model(args.data_dir, "Lite-hrnet-x-mod3")
67-
retrieve_otx_model(args.data_dir, "is_efficientnetb2b_maskrcnn_coco_reduced")
68-
retrieve_otx_model(
69-
args.data_dir, "is_efficientnetb2b_maskrcnn_coco_reduced_onnx", "onnx"
70-
)
71-
retrieve_otx_model(args.data_dir, "is_resnet50_maskrcnn_coco_reduced")
72-
retrieve_otx_model(args.data_dir, "mobilenet_v3_large_hc_cf")
73-
retrieve_otx_model(args.data_dir, "classification_model_with_xai_head")
74-
retrieve_otx_model(args.data_dir, "detection_model_with_xai_head")
75-
retrieve_otx_model(args.data_dir, "segmentation_model_with_xai_head")
76-
retrieve_otx_model(args.data_dir, "maskrcnn_model_with_xai_head")
77-
retrieve_otx_model(args.data_dir, "maskrcnn_xai_tiling")
78-
retrieve_otx_model(args.data_dir, "tile_classifier")
79-
retrieve_otx_model(args.data_dir, "anomaly_padim_bottle_mvtec")
80-
retrieve_otx_model(args.data_dir, "anomaly_stfpm_bottle_mvtec")
81-
retrieve_otx_model(args.data_dir, "deit-tiny")
63+
otx_models_dir = args.data_dir / "otx_models"
64+
otx_models_dir.mkdir(parents=True, exist_ok=True)
65+
async with httpx.AsyncClient(timeout=20.0) as client:
66+
await asyncio.gather(
67+
download_images(args.data_dir),
68+
download_otx_model(client, otx_models_dir, "mlc_mobilenetv3_large_voc"),
69+
download_otx_model(client, otx_models_dir, "mlc_efficient_b0_voc"),
70+
download_otx_model(client, otx_models_dir, "mlc_efficient_v2s_voc"),
71+
download_otx_model(client, otx_models_dir, "det_mobilenetv2_atss_bccd"),
72+
download_otx_model(
73+
client, otx_models_dir, "det_mobilenetv2_atss_bccd_onnx", "onnx"
74+
),
75+
download_otx_model(client, otx_models_dir, "cls_mobilenetv3_large_cars"),
76+
download_otx_model(
77+
client, otx_models_dir, "cls_mobilenetv3_large_cars", "onnx"
78+
),
79+
download_otx_model(client, otx_models_dir, "cls_efficient_b0_cars"),
80+
download_otx_model(client, otx_models_dir, "cls_efficient_v2s_cars"),
81+
download_otx_model(client, otx_models_dir, "Lite-hrnet-18"),
82+
download_otx_model(client, otx_models_dir, "Lite-hrnet-18_mod2"),
83+
download_otx_model(client, otx_models_dir, "Lite-hrnet-s_mod2"),
84+
download_otx_model(client, otx_models_dir, "Lite-hrnet-s_mod2", "onnx"),
85+
download_otx_model(client, otx_models_dir, "Lite-hrnet-x-mod3"),
86+
download_otx_model(
87+
client, otx_models_dir, "is_efficientnetb2b_maskrcnn_coco_reduced"
88+
),
89+
download_otx_model(
90+
client,
91+
otx_models_dir,
92+
"is_efficientnetb2b_maskrcnn_coco_reduced_onnx",
93+
"onnx",
94+
),
95+
download_otx_model(
96+
client, otx_models_dir, "is_resnet50_maskrcnn_coco_reduced"
97+
),
98+
download_otx_model(client, otx_models_dir, "mobilenet_v3_large_hc_cf"),
99+
download_otx_model(
100+
client, otx_models_dir, "classification_model_with_xai_head"
101+
),
102+
download_otx_model(client, otx_models_dir, "detection_model_with_xai_head"),
103+
download_otx_model(
104+
client, otx_models_dir, "segmentation_model_with_xai_head"
105+
),
106+
download_otx_model(client, otx_models_dir, "maskrcnn_model_with_xai_head"),
107+
download_otx_model(client, otx_models_dir, "maskrcnn_xai_tiling"),
108+
download_otx_model(client, otx_models_dir, "tile_classifier"),
109+
download_otx_model(client, otx_models_dir, "anomaly_padim_bottle_mvtec"),
110+
download_otx_model(client, otx_models_dir, "anomaly_stfpm_bottle_mvtec"),
111+
download_otx_model(client, otx_models_dir, "deit-tiny"),
112+
)
113+
114+
115+
if __name__ == "__main__":
116+
asyncio.run(main())

0 commit comments

Comments
 (0)