Update app.py

app.py

@@ -7,25 +7,11 @@ from transformers import ResNetForImageClassification, AdamW
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 import streamlit as st
-import gradio as gr
-import os
-import zipfile
-import numpy as np
-import torch
-from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor
-from transformers import ResNetForImageClassification, AdamW
-from PIL import Image
-from torch.utils.data import Dataset, DataLoader
-import streamlit as st
-import gradio as gr
 
 # Load feature extractor and model
 feature_extractor = SegformerFeatureExtractor.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512')
 segformer_model = SegformerForSemanticSegmentation.from_pretrained('nvidia/segformer-b0-finetuned-ade-512-512')
 
-
-
-
 # Function to extract zip files
 def extract_zip(zip_file, extract_to):
     with zipfile.ZipFile(zip_file, 'r') as zip_ref:
@@ -34,7 +20,7 @@ def extract_zip(zip_file, extract_to):
 # Preprocess images
 def preprocess_image(image_path):
     ext = os.path.splitext(image_path)[-1].lower()
-
+    
     if ext == '.npy':
         image_data = np.load(image_path)
         image_tensor = torch.tensor(image_data).float()
@@ -62,13 +48,19 @@ def prepare_dataset(extracted_folder):
     image_paths = []
     labels = []
     
-    for disease_folder in ['alzheimers_dataset', 'parkinsons_dataset', 'MSjpg']:
+    # Define the mapping of folders to labels
+    folder_label_mapping = {
+        'alzheimers_dataset': 0,
+        'parkinsons_dataset': 1,
+        'MSjpg': 2
+    }
+
+    for disease_folder, label in folder_label_mapping.items():
         folder_path = os.path.join(neuronii_path, disease_folder)
         
         if not os.path.exists(folder_path):
             print(f"Folder not found: {folder_path}")
             continue  
-        label = {'alzheimers_dataset': 0, 'parkinsons_dataset': 1, 'MSjpg': 2}[disease_folder]
         
         for img_file in os.listdir(folder_path):
             if img_file.endswith(('.npy', '.jpg', '.jpeg')):
@@ -97,9 +89,7 @@ class CustomImageDataset(Dataset):
 def fine_tune_classification_model(train_loader):
     # Load the ResNet model with ignore_mismatched_sizes
     model = ResNetForImageClassification.from_pretrained('microsoft/resnet-50', num_labels=3, ignore_mismatched_sizes=True)
-        # Print model architecture to identify the classifier layer
-    print(model)  # Inspect the model structure
-
+    
     # Update the classifier layer to match the number of labels
     if hasattr(model, 'classifier'):
         if isinstance(model.classifier, torch.nn.Sequential):
@@ -116,7 +106,7 @@ def fine_tune_classification_model(train_loader):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model.to(device)
     
-    for epoch in range(10):
+    for epoch in range(10):  # Adjust epochs as needed
         running_loss = 0.0
         for images, labels in train_loader:
             images, labels = images.to(device), labels.to(device)
@@ -164,7 +154,7 @@ def fine_tune_segmentation_model(train_loader):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model.to(device)
     
-    for epoch in range(10):
+    for epoch in range(10):  # Adjust epochs as needed
         running_loss = 0.0
         for images, labels in train_loader:
             images, labels = images.to(device), labels.to(device)
@@ -185,3 +175,4 @@ if st.button("Start Segmentation Training"):
     # Fine-tune the segmentation model
     final_loss_seg = fine_tune_segmentation_model(seg_train_loader)
     st.write(f"Segmentation Training Complete with Final Loss: {final_loss_seg}")
+