README updated with the dataset loding.

README.md

@@ -104,6 +104,25 @@ config_dict = {
 # Initialize model
 model = PCamClassifier(config_dict)
 model.load_state_dict(torch.load('best_enhanced_pcam_model.pt'))
+
+
+class PCamDataset(Dataset):
+    def __init__(self, dataset):
+        self.dataset = dataset
+        
+    def __len__(self):
+        return len(self.dataset)
+        
+    def __getitem__(self, idx):
+        example = self.dataset[idx]
+        image = example["image"].convert("RGB")
+        image_array = np.array(image) / 255.0
+        image_array = image_array.transpose(2, 0, 1).astype(np.float32)
+        return {
+            "pixel_values": image_array,
+            "labels": example["label"]
+        }
+
 ```
 
 ---