Dynamic subsampling

app.py

@@ -63,9 +63,45 @@ print(f"Loaded {len(images_data)} images for the quiz")
 with open('images.txt', 'r') as f:
     image_filenames = [line.strip() for line in f.readlines() if line.strip()]
 
-# Initialize CODA
+# Initialize CODA with subsampled dataset
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-dataset = Dataset("iwildcam_demo.pt", device=device)
+
+# Load full dataset
+full_preds = torch.load("iwildcam_demo.pt").to(device)
+full_labels = torch.load("iwildcam_demo_labels.pt").to(device)
+
+# Subsample to balance classes
+from collections import defaultdict
+class_to_indices = defaultdict(list)
+for idx, label in enumerate(full_labels):
+    class_idx = label.item()
+    class_to_indices[class_idx].append(idx)
+
+# Find minimum class size
+min_class_size = min(len(indices) for indices in class_to_indices.values())
+print(f"Subsampling to {min_class_size} images per class (total: {min_class_size * len(class_to_indices)} images)")
+
+# Randomly subsample each class
+subsampled_indices = []
+for class_idx in sorted(class_to_indices.keys()):
+    indices = class_to_indices[class_idx]
+    sampled = np.random.choice(indices, size=min_class_size, replace=False)
+    subsampled_indices.extend(sampled.tolist())
+
+# Sort indices to maintain order
+subsampled_indices.sort()
+
+# Create subsampled dataset
+subsampled_preds = full_preds[:, subsampled_indices, :]
+subsampled_labels = full_labels[subsampled_indices]
+image_filenames = [image_filenames[idx] for idx in subsampled_indices]
+
+# Create Dataset object with subsampled data
+dataset = Dataset.__new__(Dataset)
+dataset.preds = subsampled_preds
+dataset.labels = subsampled_labels
+dataset.device = device
+
 loss_fn = LOSS_FNS['acc']
 oracle = Oracle(dataset, loss_fn=loss_fn)