Update traininginVIT

traininginVIT

@@ -1,6 +1,11 @@
 from torch.utils.data import DataLoader, Dataset
 import torch
 from transformers import ViTForImageClassification, AdamW
+import os
+import numpy as np
+import torch
+import streamlit as st
+from transformers import ViTForImageClassification, ViTImageProcessor
 
 # Custom dataset class for loading images
 class MRIDataset(Dataset):
@@ -54,3 +59,21 @@ for epoch in range(num_epochs):
 
 # Save the fine-tuned model
 torch.save(model.state_dict(), 'vit_finetuned.pth')
+
+def fine_tune_model():
+    # Your fine-tuning logic goes here (using the ViT model)
+    num_epochs = 10
+    running_loss = 0.0
+    for epoch in range(num_epochs):
+        # Fine-tuning loop (train the model)
+        # ...
+        running_loss += 0.5  # Just a placeholder for demo purposes
+    return running_loss  # Return the final loss after training
+
+# Streamlit UI to trigger fine-tuning and display results
+st.title("MRI Image Fine-Tuning with ViT")
+
+if st.button("Start Training"):
+    # Run the fine-tuning loop when the button is clicked
+    final_loss = fine_tune_model()  # Call the function where your fine-tuning loop is
+    st.write(f"Training complete with final loss: {final_loss}")