Create app.py

app.py

@@ -0,0 +1,48 @@
+import os
+import streamlit as st
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+from preprocessing import preprocess_image
+from train import fine_tune_model
+
+# Define a function to load and preprocess images
+def load_images(image_paths):
+    images = []
+    labels = []  # Assuming you have labels, adjust as needed
+
+    for path in image_paths:
+        # Preprocess each image and append to the list
+        image_tensor = preprocess_image(path)
+        images.append(image_tensor)
+
+        # Append corresponding label (this is just an example; adjust as needed)
+        label = ...  # Replace with logic to get the label for the image
+        labels.append(label)
+
+    # Stack images into a single tensor and create a DataLoader
+    images_tensor = torch.stack(images)  # Stack into a single tensor
+    labels_tensor = torch.tensor(labels)  # Convert labels to tensor
+    dataset = TensorDataset(images_tensor, labels_tensor)
+    return DataLoader(dataset, batch_size=16, shuffle=True)  # Adjust batch_size as needed
+
+# Streamlit UI
+st.title("MRI Image Fine-Tuning with ViT")
+
+# Upload images or provide a directory path
+image_directory = st.text_input("Enter the directory of images:")
+
+if st.button("Load Images"):
+    # Get all image paths from the directory
+    if image_directory and os.path.isdir(image_directory):
+        image_paths = [os.path.join(image_directory, f) for f in os.listdir(image_directory) if f.endswith(('.nii', '.jpg', '.jpeg'))]
+        st.success(f"Loaded {len(image_paths)} images.")
+
+        # Preprocess images and create DataLoader
+        train_loader = load_images(image_paths)
+
+        # Button to start training
+        if st.button("Start Training"):
+            final_loss = fine_tune_model(train_loader)  # Trigger fine-tuning
+            st.write(f"Training complete with final loss: {final_loss}")
+    else:
+        st.error("Please enter a valid directory.")