Update app.py

app.py

@@ -3,7 +3,6 @@ import zipfile
 import numpy as np
 import torch
 from transformers import ViTForImageClassification, AdamW
-import nibabel as nib
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 import streamlit as st
@@ -17,13 +16,14 @@ def extract_zip(zip_file, extract_to):
 def preprocess_image(image_path):
     ext = os.path.splitext(image_path)[-1].lower()
 
-    if ext in ['.nii', '.nii.gz']:
-        nii_image = nib.load(image_path)
-        image_data = nii_image.get_fdata()
+    if ext in ['.npy']:
+        image_data = np.load(image_path)
         image_tensor = torch.tensor(image_data).float()
-        if len(image_tensor.shape) == 3:
-            image_tensor = image_tensor.unsqueeze(0)
-
+        if len(image_tensor.shape) == 2:  # If the image is 2D (grayscale)
+            image_tensor = image_tensor.unsqueeze(0)  # Add channel dimension
+        elif len(image_tensor.shape) == 3:  # If the image is 3D (height, width, channels)
+            image_tensor = image_tensor.permute(2, 0, 1).float()  # Change to (C, H, W)
+    
     elif ext in ['.jpg', '.jpeg']:
         img = Image.open(image_path).convert('RGB').resize((224, 224))
         img_np = np.array(img)
@@ -52,16 +52,18 @@ def prepare_dataset(extracted_folder):
         # Check if the subfolder exists
         if not os.path.exists(folder_path):
             print(f"Folder not found: {folder_path}")
-            continue  # Skip this folder if it's not foun
+            continue  # Skip this folder if it's not found
+            
         label = {'alzheimers_dataset': 0, 'parkinsons_dataset': 1, 'MSjpg': 2}[disease_folder]
         
         for img_file in os.listdir(folder_path):
-            if img_file.endswith(('.nii', '.jpg', '.jpeg')):
+            if img_file.endswith(('.npy', '.jpg', '.jpeg')):
                 image_paths.append(os.path.join(folder_path, img_file))
                 labels.append(label)
             else:
-                print(f"Unsuported file:{img_file}")
-    print(f"Total images loaded :{len(image_paths)}")
+                print(f"Unsupported file: {img_file}")
+    
+    print(f"Total images loaded: {len(image_paths)}")
     return image_paths, labels
 
 # Custom Dataset class