Process images as numpy arrays

README.md

@@ -92,16 +92,17 @@ Given n samples, the value of the Vendi Score ranges between 1 and n, with highe
 >>> samples = [0, 0, 10, 10, 20, 20]
 >>> k = lambda a, b: np.exp(-np.abs(a - b))
 >>> vendiscore.compute(samples=samples, k=k)
-{"VS": 2.9999...}
+{'VS': 2.9999...}
 ```
 
 If you already have precomputed a similarity matrix:
 ```
+>>> vendiscore = evaluate.load("danf0/vendiscore", "K")
 >>> K = np.array([[1.0, 0.9, 0.0],
                   [0.9, 1.0, 0.0],
                   [0.0, 0.0, 1.0]])
 >>> vendiscore.compute(samples=K, score_K=True)
-2.1573
+{'VS': 2.1573...}
 ```
 
 If your similarity function is a dot product between `n` normalized
@@ -109,9 +110,10 @@ If your similarity function is a dot product between `n` normalized
 to compute the Vendi Score using the covariance matrix, `X @ X.T`.
 (If the rows of `X` are not normalized, set `normalize = True`.)
 ```
->>> X = np.array([[100, 0], [99, 1], [1, 99], [0, 100])
+>>> vendiscore = evaluate.load("danf0/vendiscore", "X")
+>>> X = np.array([[100, 0], [99, 1], [1, 99], [0, 100]])
 >>> vendiscore.compute(samples=X, score_dual=True, normalize=True)
-1.9989...
+{'VS': 1.99989...}
 ```
 
 Image similarity can be calculated using inner products between pixel vectors or between embeddings from a neural network.
@@ -137,16 +139,12 @@ The default embeddings are from the pool-2048 layer of the torchvision version o
 
 Text similarity can be calculated using n-gram overlap or using inner products between embeddings from a neural network.
 ```
->>> sents = ["Look, Jane.",
-             "See Spot.",
-             "See Spot run.",
-             "Run, Spot, run.",
-	     "Jane sees Spot run."]
->>> ngram_vs = vendiscore.compute(samples=sents, k="ngram_overlap", ns=[1, 2])
->>> bert_vs = vendiscore.compute(samples=sents, k="text_embeddings", model_path="bert-base-uncased")
->>> simcse_vs = vendiscore.compute(samples=sents, k="text_embeddings", model_path="princeton-nlp/unsup-simcse-bert-base-uncased")
->>> print(f"N-grams: {ngram_vs:.02f}, BERT: {bert_vs:.02f}, SimCSE: {simcse_vs:.02f})
-N-grams: 3.91, BERT: 1.21, SimCSE: 2.81
+>>> vendiscore = evaluate.load("danf0/vendiscore", "text")
+>>> sents = ["Look, Jane.", "See Spot.", "See Spot run.", "Run, Spot, run.", "Jane sees Spot run."]
+>>> ngram_vs = vendiscore.compute(samples=sents, k="ngram_overlap", ns=[1, 2])["VS"]
+>>> bert_vs = vendiscore.compute(samples=sents, k="text_embeddings", model_path="bert-base-uncased")["VS"]
+>>> print(f"N-grams: {ngram_vs:.02f}, BERT: {bert_vs:.02f}")
+N-grams: 3.91, BERT: 1.21
 ```
 
 ## Limitations and Bias

vendiscore.py

@@ -14,6 +14,8 @@
 import evaluate
 import datasets
 import numpy as np
+import PIL
+from PIL import Image
 
 from vendi_score import vendi, image_utils, text_utils
 
@@ -69,24 +71,11 @@ Examples:
 """
 
 
-def get_dtype(config_name):
-    if config_name == "text":
-        return datasets.Features({"samples": datasets.Value("string")})
-    if config_name == "image":
-        return datasets.Features({"samples": datasets.Image})
-    elif config_name in ("X", "K"):
-        return datasets.Array2D
-    elif config_name == "default":
-        return datasets.Value("string")
-    else:
-        return datasets.Value(config_name)
-
-
 def get_features(config_name):
     if config_name in ("text", "default"):
         return datasets.Features({"samples": datasets.Value("string")})
     if config_name == "image":
-        return datasets.Features({"samples": datasets.Image})
+        return datasets.Features({"samples": datasets.Array3D})
     if config_name in ("K", "X"):
         return [
             datasets.Features(
@@ -164,10 +153,12 @@ class VendiScore(evaluate.Metric):
                 model_path=model_path,
             )
         elif type(k) == str and k == "pixels":
-            vs = image_utils.pixel_vendi_score(samples)
+            vs = image_utils.pixel_vendi_score(
+                [Image.fromarray(x) for x in samples]
+            )
         elif type(k) == str and k == "image_embeddings":
             vs = image_utils.embedding_vendi_score(
-                samples,
+                [Image.fromarray(x) for x in samples],
                 batch_size=batch_size,
                 device=device,
                 model=model,