Update app.py

app.py

@@ -1,48 +1,11 @@
+import zipfile
 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
+def extract_zip(zip_file, extract_to):
+    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+        zip_ref.extractall(extract_to)
 
-    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.")
+# Example usage
+zip_file_path = 'path_to_your_zip_file.zip'
+extract_to = 'path_to_extracted_images/'
+extract_zip(zip_file_path, extract_to)