Update README.md

README.md

@@ -28,9 +28,11 @@ To load a pre-trained model ("VisionTransformer.pt"), use the following code sni
 
 ```python
 import clip
-from clip.downstream_task import TaskType
 import torch  # Make sure to import torch
 
+from clip.downstream_task import TaskType
+
+
 device = "cpu"  # Change to 'cuda' if you have a GPU
 num_classes = 4  # Number of classes in the original HYPERVIEW dataset
 
@@ -43,4 +45,28 @@ model, _ = clip.load(
 # Load the pre-trained weights
 model.load_state_dict(torch.load("VisionTransformer.pt"))
 model.eval()  # Set the model to evaluation mode
-```
+```
+
+### Loading training data
+
+```python
+import numpy as np
+
+from clip.hyperview_data_loader import HyperDataloader, DataReader
+
+
+im_size = 224  # Image size
+num_classes = 4  # Number of classes in the original HYPERVIEW dataset
+
+# Paths to training data and ground truth
+train_path = "<TRAIN_PATH>"
+train_gt_path = "<TRAIN_PATH>/train_gt.csv"
+
+# Initialize the dataset reader and transformations
+target_index = list(np.arange(num_classes))
+trans_tr, _ = HyperDataloader._init_transform(im_size)
+train_dataset = DataReader(
+    database_dir=train_path, label_paths=train_gt_path,
+    transform=trans_tr, target_index=target_index
+)
+````