You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
210 lines
11 KiB
210 lines
11 KiB
import sys
|
|
import json
|
|
import numpy as np
|
|
from PIL import Image
|
|
from pathlib import Path
|
|
import boto3, botocore
|
|
from tinygrad import Tensor, dtypes
|
|
from tinygrad.helpers import fetch, tqdm, getenv
|
|
import pandas as pd
|
|
import concurrent.futures
|
|
|
|
BASEDIR = Path(__file__).parent / "open-images-v6-mlperf"
|
|
BUCKET_NAME = "open-images-dataset"
|
|
TRAIN_BBOX_ANNOTATIONS_URL = "https://storage.googleapis.com/openimages/v6/oidv6-train-annotations-bbox.csv"
|
|
VALIDATION_BBOX_ANNOTATIONS_URL = "https://storage.googleapis.com/openimages/v5/validation-annotations-bbox.csv"
|
|
MAP_CLASSES_URL = "https://storage.googleapis.com/openimages/v5/class-descriptions-boxable.csv"
|
|
MLPERF_CLASSES = ['Airplane', 'Antelope', 'Apple', 'Backpack', 'Balloon', 'Banana',
|
|
'Barrel', 'Baseball bat', 'Baseball glove', 'Bee', 'Beer', 'Bench', 'Bicycle',
|
|
'Bicycle helmet', 'Bicycle wheel', 'Billboard', 'Book', 'Bookcase', 'Boot',
|
|
'Bottle', 'Bowl', 'Bowling equipment', 'Box', 'Boy', 'Brassiere', 'Bread',
|
|
'Broccoli', 'Bronze sculpture', 'Bull', 'Bus', 'Bust', 'Butterfly', 'Cabinetry',
|
|
'Cake', 'Camel', 'Camera', 'Candle', 'Candy', 'Cannon', 'Canoe', 'Carrot', 'Cart',
|
|
'Castle', 'Cat', 'Cattle', 'Cello', 'Chair', 'Cheese', 'Chest of drawers', 'Chicken',
|
|
'Christmas tree', 'Coat', 'Cocktail', 'Coffee', 'Coffee cup', 'Coffee table', 'Coin',
|
|
'Common sunflower', 'Computer keyboard', 'Computer monitor', 'Convenience store',
|
|
'Cookie', 'Countertop', 'Cowboy hat', 'Crab', 'Crocodile', 'Cucumber', 'Cupboard',
|
|
'Curtain', 'Deer', 'Desk', 'Dinosaur', 'Dog', 'Doll', 'Dolphin', 'Door', 'Dragonfly',
|
|
'Drawer', 'Dress', 'Drum', 'Duck', 'Eagle', 'Earrings', 'Egg (Food)', 'Elephant',
|
|
'Falcon', 'Fedora', 'Flag', 'Flowerpot', 'Football', 'Football helmet', 'Fork',
|
|
'Fountain', 'French fries', 'French horn', 'Frog', 'Giraffe', 'Girl', 'Glasses',
|
|
'Goat', 'Goggles', 'Goldfish', 'Gondola', 'Goose', 'Grape', 'Grapefruit', 'Guitar',
|
|
'Hamburger', 'Handbag', 'Harbor seal', 'Headphones', 'Helicopter', 'High heels',
|
|
'Hiking equipment', 'Horse', 'House', 'Houseplant', 'Human arm', 'Human beard',
|
|
'Human body', 'Human ear', 'Human eye', 'Human face', 'Human foot', 'Human hair',
|
|
'Human hand', 'Human head', 'Human leg', 'Human mouth', 'Human nose', 'Ice cream',
|
|
'Jacket', 'Jeans', 'Jellyfish', 'Juice', 'Kitchen & dining room table', 'Kite',
|
|
'Lamp', 'Lantern', 'Laptop', 'Lavender (Plant)', 'Lemon', 'Light bulb', 'Lighthouse',
|
|
'Lily', 'Lion', 'Lipstick', 'Lizard', 'Man', 'Maple', 'Microphone', 'Mirror',
|
|
'Mixing bowl', 'Mobile phone', 'Monkey', 'Motorcycle', 'Muffin', 'Mug', 'Mule',
|
|
'Mushroom', 'Musical keyboard', 'Necklace', 'Nightstand', 'Office building',
|
|
'Orange', 'Owl', 'Oyster', 'Paddle', 'Palm tree', 'Parachute', 'Parrot', 'Pen',
|
|
'Penguin', 'Personal flotation device', 'Piano', 'Picture frame', 'Pig', 'Pillow',
|
|
'Pizza', 'Plate', 'Platter', 'Porch', 'Poster', 'Pumpkin', 'Rabbit', 'Rifle',
|
|
'Roller skates', 'Rose', 'Salad', 'Sandal', 'Saucer', 'Saxophone', 'Scarf', 'Sea lion',
|
|
'Sea turtle', 'Sheep', 'Shelf', 'Shirt', 'Shorts', 'Shrimp', 'Sink', 'Skateboard',
|
|
'Ski', 'Skull', 'Skyscraper', 'Snake', 'Sock', 'Sofa bed', 'Sparrow', 'Spider', 'Spoon',
|
|
'Sports uniform', 'Squirrel', 'Stairs', 'Stool', 'Strawberry', 'Street light',
|
|
'Studio couch', 'Suit', 'Sun hat', 'Sunglasses', 'Surfboard', 'Sushi', 'Swan',
|
|
'Swimming pool', 'Swimwear', 'Tank', 'Tap', 'Taxi', 'Tea', 'Teddy bear', 'Television',
|
|
'Tent', 'Tie', 'Tiger', 'Tin can', 'Tire', 'Toilet', 'Tomato', 'Tortoise', 'Tower',
|
|
'Traffic light', 'Train', 'Tripod', 'Truck', 'Trumpet', 'Umbrella', 'Van', 'Vase',
|
|
'Vehicle registration plate', 'Violin', 'Wall clock', 'Waste container', 'Watch',
|
|
'Whale', 'Wheel', 'Wheelchair', 'Whiteboard', 'Window', 'Wine', 'Wine glass', 'Woman',
|
|
'Zebra', 'Zucchini',
|
|
]
|
|
|
|
|
|
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}")
|
|
|
|
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
|
|
def extract_dims(path): return Image.open(path).size[::-1]
|
|
|
|
def export_to_coco(class_map, annotations, image_list, dataset_path, output_path, subset, classes=MLPERF_CLASSES):
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
cats = [{"id": i, "name": c, "supercategory": None} for i, c in enumerate(classes)]
|
|
categories_map = pd.DataFrame([(i, c) for i, c in enumerate(classes)], columns=["category_id", "category_name"])
|
|
class_map = class_map.merge(categories_map, left_on="DisplayName", right_on="category_name", how="inner")
|
|
annotations = annotations[annotations["ImageID"].isin(image_list)]
|
|
annotations = annotations.merge(class_map, on="LabelName", how="inner")
|
|
annotations["image_id"] = pd.factorize(annotations["ImageID"].tolist())[0]
|
|
annotations[["height", "width"]] = annotations.apply(lambda x: extract_dims(dataset_path / f"{x['ImageID']}.jpg"), axis=1, result_type="expand")
|
|
|
|
# Images
|
|
imgs = [{"id": int(id + 1), "file_name": f"{image_id}.jpg", "height": row["height"], "width": row["width"], "subset": subset, "license": None, "coco_url": None}
|
|
for (id, image_id), row in (annotations.groupby(["image_id", "ImageID"]).first().iterrows())
|
|
]
|
|
|
|
# Annotations
|
|
annots = []
|
|
for i, row in annotations.iterrows():
|
|
xmin, ymin, xmax, ymax, img_w, img_h = [row[k] for k in ["XMin", "YMin", "XMax", "YMax", "width", "height"]]
|
|
x, y, w, h = xmin * img_w, ymin * img_h, (xmax - xmin) * img_w, (ymax - ymin) * img_h
|
|
coco_annot = {"id": int(i) + 1, "image_id": int(row["image_id"] + 1), "category_id": int(row["category_id"]), "bbox": [x, y, w, h], "area": w * h}
|
|
coco_annot.update({k: row[k] for k in ["IsOccluded", "IsInside", "IsDepiction", "IsTruncated", "IsGroupOf"]})
|
|
coco_annot["iscrowd"] = int(row["IsGroupOf"])
|
|
annots.append(coco_annot)
|
|
|
|
info = {"dataset": "openimages_mlperf", "version": "v6"}
|
|
coco_annotations = {"info": info, "licenses": [], "categories": cats, "images": imgs, "annotations": annots}
|
|
with open(output_path, "w") as fp:
|
|
json.dump(coco_annotations, fp)
|
|
|
|
def get_image_list(class_map, annotations, classes=MLPERF_CLASSES):
|
|
labels = class_map[class_map["DisplayName"].isin(classes)]["LabelName"]
|
|
image_ids = annotations[annotations["LabelName"].isin(labels)]["ImageID"].unique()
|
|
return image_ids
|
|
|
|
def download_image(bucket, subset, image_id, data_dir):
|
|
try:
|
|
bucket.download_file(f"{subset}/{image_id}.jpg", f"{data_dir}/{image_id}.jpg")
|
|
except botocore.exceptions.ClientError as exception:
|
|
sys.exit(f"ERROR when downloading image `validation/{image_id}`: {str(exception)}")
|
|
|
|
def fetch_openimages(output_fn:str, base_dir:Path, subset:str):
|
|
bucket = boto3.resource("s3", config=botocore.config.Config(signature_version=botocore.UNSIGNED)).Bucket(BUCKET_NAME)
|
|
|
|
annotations_dir, data_dir = base_dir / "annotations", base_dir / f"{subset}/data"
|
|
annotations_dir.mkdir(parents=True, exist_ok=True)
|
|
data_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
if subset == "train":
|
|
annotations_fn = annotations_dir / TRAIN_BBOX_ANNOTATIONS_URL.split('/')[-1]
|
|
fetch(TRAIN_BBOX_ANNOTATIONS_URL, annotations_fn)
|
|
else: # subset == validation
|
|
annotations_fn = annotations_dir / VALIDATION_BBOX_ANNOTATIONS_URL.split('/')[-1]
|
|
fetch(VALIDATION_BBOX_ANNOTATIONS_URL, annotations_fn)
|
|
|
|
annotations = pd.read_csv(annotations_fn)
|
|
|
|
classmap_fn = annotations_dir / MAP_CLASSES_URL.split('/')[-1]
|
|
fetch(MAP_CLASSES_URL, classmap_fn)
|
|
class_map = pd.read_csv(classmap_fn, names=["LabelName", "DisplayName"])
|
|
|
|
image_list = get_image_list(class_map, annotations)
|
|
|
|
with concurrent.futures.ThreadPoolExecutor() as executor:
|
|
futures = [executor.submit(download_image, bucket, subset, image_id, data_dir) for image_id in image_list]
|
|
for future in (t := tqdm(concurrent.futures.as_completed(futures), total=len(image_list))):
|
|
t.set_description(f"Downloading images")
|
|
future.result()
|
|
|
|
print("Converting annotations to COCO format...")
|
|
export_to_coco(class_map, annotations, image_list, data_dir, output_fn, subset)
|
|
|
|
def image_load(base_dir, subset, fn):
|
|
img_folder = base_dir / f"{subset}/data"
|
|
return Image.open(img_folder / fn).convert('RGB')
|
|
|
|
def prepare_target(annotations, img_id, img_size):
|
|
boxes = [annot["bbox"] for annot in annotations]
|
|
boxes = np.array(boxes, dtype=np.float32).reshape(-1, 4)
|
|
boxes[:, 2:] += boxes[:, :2]
|
|
boxes[:, 0::2] = boxes[:, 0::2].clip(0, img_size[1])
|
|
boxes[:, 1::2] = boxes[:, 1::2].clip(0, img_size[0])
|
|
keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
|
|
boxes = boxes[keep]
|
|
classes = [annot["category_id"] for annot in annotations]
|
|
classes = np.array(classes, dtype=np.int64)
|
|
classes = classes[keep]
|
|
return {"boxes": boxes, "labels": classes, "image_id": img_id, "image_size": img_size}
|
|
|
|
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 = resize(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) -> 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, ann_file)
|
|
print("Done")
|
|
|
|
return ann_file
|
|
|
|
def random_horizontal_flip(img, tgt, prob=0.5):
|
|
import torch
|
|
import torchvision.transforms.functional as F
|
|
if torch.rand(1) < prob:
|
|
w = img.size[0]
|
|
img = F.hflip(img)
|
|
tgt["boxes"][:, [0, 2]] = w - tgt["boxes"][:, [2, 0]]
|
|
return img, tgt
|
|
|
|
def resize(img:Image, tgt:dict[str, np.ndarray|tuple]|None=None, size:tuple[int, int]=(800, 800)) -> tuple[np.ndarray, np.ndarray, tuple]|tuple[np.ndarray, tuple]:
|
|
import torchvision.transforms.functional as F
|
|
img_size = img.size[::-1]
|
|
img = F.resize(img, size=size)
|
|
img = np.array(img)
|
|
|
|
if tgt is not None:
|
|
ratios = [s / s_orig for s, s_orig in zip(size, img_size)]
|
|
ratio_h, ratio_w = ratios
|
|
x_min, y_min, x_max, y_max = [tgt["boxes"][:, i] for i in range(tgt["boxes"].shape[-1])]
|
|
x_min = x_min * ratio_w
|
|
x_max = x_max * ratio_w
|
|
y_min = y_min * ratio_h
|
|
y_max = y_max * ratio_h
|
|
|
|
tgt["boxes"] = np.stack([x_min, y_min, x_max, y_max], axis=1)
|
|
return img, tgt, img_size
|
|
|
|
return img, img_size
|
|
|
|
if __name__ == "__main__":
|
|
download_dataset(base_dir:=getenv("BASE_DIR", BASEDIR), "train")
|
|
download_dataset(base_dir, "validation")
|
|
|