up

app.py

@@ -1,5 +1,5 @@
 from datasets import load_dataset
-from collections import Counter
+from collections import Counter, defaultdict
 from random import sample, shuffle
 import datasets
 from pandas import DataFrame
@@ -65,10 +65,19 @@ def start():
 
     # sort by count
     ids = sorted(ids.items(), key=lambda x: x[1])
-    ids = [i[0] for i in ids]
+    freq_ids = defaultdict(list)
+    for k, v in ids:
+        freq_ids[v].append(k)
+
+    # shuffle in-between categories
+    for k, v_list in freq_ids.items():
+        shuffle(v_list)
+        freq_ids[v] = v_list
+
+    shuffled_ids = sum(list(freq_ids.values()), [])
 
     # get lowest count ids
-    id_candidates = ids[: (10 * NUM_QUESTIONS)]
+    id_candidates = shuffled_ids[: (10 * NUM_QUESTIONS)]
 
     # get random `NUM_QUESTIONS` ids to check
     image_ids = sample(id_candidates, k=NUM_QUESTIONS)