diff --git a/README.md b/README.md index a103d80f..35178b90 100644 --- a/README.md +++ b/README.md @@ -71,6 +71,8 @@ We have a number of precomputed data sets in HDF5 format. All data sets have bee | [NYTimes](https://archive.ics.uci.edu/ml/datasets/bag+of+words) | 256 | 290,000 | 10,000 | 100 | Angular | [HDF5](http://ann-benchmarks.com/nytimes-256-angular.hdf5) (301MB) | | [SIFT](http://corpus-texmex.irisa.fr/) | 128 | 1,000,000 | 10,000 | 100 | Euclidean | [HDF5](http://ann-benchmarks.com/sift-128-euclidean.hdf5) (501MB) | | [Last.fm](https://github.com/erikbern/ann-benchmarks/pull/91) | 65 | 292,385 | 50,000 | 100 | Angular | [HDF5](http://ann-benchmarks.com/lastfm-64-dot.hdf5) (135MB) | +| [COCO-I2I](https://cocodataset.org/) | 512 | 113,287 | 10,000 | 100 | Angular | [HDF5](https://github.com/fabiocarrara/str-encoders/releases/download/v0.1.3/coco-i2i-512-angular.hdf5) (136MB) | +| [COCO-T2I](https://cocodataset.org/) | 512 | 113,287 | 10,000 | 100 | Angular | [HDF5](https://github.com/fabiocarrara/str-encoders/releases/download/v0.1.3/coco-t2i-512-angular.hdf5) (136MB) | Results ======= diff --git a/ann_benchmarks/datasets.py b/ann_benchmarks/datasets.py index 44eb8a46..8b327878 100644 --- a/ann_benchmarks/datasets.py +++ b/ann_benchmarks/datasets.py @@ -577,6 +577,25 @@ def dbpedia_entities_openai_1M(out_fn, n = None): write_output(X_train, X_test, out_fn, "angular") +def coco(out_fn: str, kind: str): + assert kind in ('t2i', 'i2i') + + local_fn = "coco-clip-b16-512-features.hdf5" + url = "https://github.com/fabiocarrara/str-encoders/releases/download/v0.1.3/%s" % local_fn + download(url, local_fn) + + with h5py.File(local_fn, "r") as f: + img_X = f['img_feats'][:] + + X_train, X_test = train_test_split(img_X, test_size=10_000) + + if kind == 't2i': + # there are 5 captions per image, take the first one + txt_X = f['txt_feats'][::5] + _, X_test = train_test_split(txt_X, test_size=10_000) + + write_output(X_train, X_test, out_fn, "angular") + DATASETS: Dict[str, Callable[[str], None]] = { "deep-image-96-angular": deep_image, @@ -606,6 +625,8 @@ def dbpedia_entities_openai_1M(out_fn, n = None): "movielens1m-jaccard": movielens1m, "movielens10m-jaccard": movielens10m, "movielens20m-jaccard": movielens20m, + "coco-i2i-512-angular": lambda out_fn: coco(out_fn, "i2i"), + "coco-t2i-512-angular": lambda out_fn: coco(out_fn, "t2i"), } DATASETS.update({