diff --git a/examples/mlperf/model_eval.py b/examples/mlperf/model_eval.py index cef6d6cd37..9d828d07ca 100644 --- a/examples/mlperf/model_eval.py +++ b/examples/mlperf/model_eval.py @@ -94,11 +94,11 @@ def eval_retinanet(): x /= input_std return x - from extra.datasets.openimages import openimages, iterate + from extra.datasets.openimages import download_dataset, iterate, BASEDIR from pycocotools.coco import COCO from pycocotools.cocoeval import COCOeval from contextlib import redirect_stdout - coco = COCO(openimages('validation')) + coco = COCO(download_dataset(base_dir:=getenv("BASE_DIR", BASEDIR), 'validation')) coco_eval = COCOeval(coco, iouType="bbox") coco_evalimgs, evaluated_imgs, ncats, narea = [], [], len(coco_eval.params.catIds), len(coco_eval.params.areaRng) @@ -107,7 +107,7 @@ def eval_retinanet(): n, bs = 0, 8 st = time.perf_counter() - for x, targets in iterate(coco, bs): + for x, targets in iterate(coco, base_dir, bs): dat = Tensor(x.astype(np.float32)) mt = time.perf_counter() if dat.shape[0] == bs: diff --git a/extra/datasets/openimages.py b/extra/datasets/openimages.py index 2e86f36ca9..52ae27829d 100644 --- a/extra/datasets/openimages.py +++ b/extra/datasets/openimages.py @@ -54,17 +54,12 @@ MLPERF_CLASSES = ['Airplane', 'Antelope', 'Apple', 'Backpack', 'Balloon', 'Banan ] -def openimages(base_dir:Path, subset:str): +def openimages(base_dir:Path, subset:str, ann_file:Path): valid_subsets = ['train', 'validation'] if subset not in valid_subsets: raise ValueError(f"{subset=} must be one of {valid_subsets}") - ann_file = base_dir / f"{subset}/labels/openimages-mlperf.json" - - if not ann_file.is_file(): - fetch_openimages(ann_file, base_dir, subset) - - return ann_file + fetch_openimages(ann_file, base_dir, subset) # this slows down the conversion a lot! # maybe use https://raw.githubusercontent.com/scardine/image_size/master/get_image_size.py @@ -142,8 +137,8 @@ def fetch_openimages(output_fn:str, base_dir:Path, subset:str): print("Converting annotations to COCO format...") export_to_coco(class_map, annotations, image_list, data_dir, output_fn, subset) -def image_load(subset, fn): - img_folder = BASEDIR / f"{subset}/data" +def image_load(base_dir, subset, fn): + img_folder = base_dir / f"{subset}/data" img = Image.open(img_folder / fn).convert('RGB') import torchvision.transforms.functional as F ret = F.resize(img, size=(800, 800)) @@ -163,26 +158,26 @@ def prepare_target(annotations, img_id, img_size): classes = classes[keep] return {"boxes": boxes, "labels": classes, "image_id": img_id, "image_size": img_size} -def iterate(coco, bs=8): +def iterate(coco, base_dir, bs=8): image_ids = sorted(coco.imgs.keys()) for i in range(0, len(image_ids), bs): X, targets = [], [] for img_id in image_ids[i:i+bs]: img_dict = coco.loadImgs(img_id)[0] - x, original_size = image_load(img_dict['subset'], img_dict["file_name"]) + x, original_size = image_load(base_dir, img_dict['subset'], img_dict["file_name"]) X.append(x) annotations = coco.loadAnns(coco.getAnnIds(img_id)) targets.append(prepare_target(annotations, img_id, original_size)) yield np.array(X), targets -def download_dataset(base_dir:Path, subset:str): - try: - (base_dir / subset).mkdir(parents=True) +def download_dataset(base_dir:Path, subset:str) -> Path: + if (ann_file:=base_dir / f"{subset}/labels/openimages-mlperf.json").is_file(): print(f"{subset} dataset is already available") + else: print(f"Downloading {subset} dataset...") - openimages(base_dir, subset) + openimages(base_dir, subset, ann_file) print("Done") - except OSError: - print(f"{subset} dataset is already available") + + return ann_file if __name__ == "__main__":