Update inference_vit.py

inference_vit.py

@@ -116,10 +116,13 @@ def load_image_by_id2(image_id):
     """
     Load image and label directly from Hugging Face dataset
     """
+    global dataset   # make sure dataset is global
+
     try:
-        img_obj = dataset[image_id]["image"]   # returns PIL.Image
-        label_idx = dataset[image_id]["label"] # integer
-        label_name = dataset.features["label"].names[label_idx]  # human-readable
+        # select train split
+        img_obj = dataset['train'][image_id]["image"]   # returns PIL.Image
+        label_idx = dataset['train'][image_id]["label"] # integer
+        label_name = dataset['train'].features["label"].names[label_idx]  # human-readable
         return img_obj, label_name
     except Exception as e:
         print(f"Error loading image {image_id}: {e}")