Update app.py

app.py

@@ -70,6 +70,7 @@ def fine_tune_model(train_loader):
     criterion = torch.nn.CrossEntropyLoss()
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model.to(device)
+    
     for epoch in range(10):
         running_loss = 0.0
         for images, labels in train_loader:
@@ -85,10 +86,25 @@ def fine_tune_model(train_loader):
 # Streamlit UI
 st.title("Fine-tune ViT on MRI Scans")
 
+# Input for zip file paths
+zip_file_1 = st.text_input("https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/archive%20(5).zip")
+zip_file_2 = st.text_input("https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/MS.zip")
+
 if st.button("Start Training"):
-    extract_zip('https://huggingface.co/spaces/Tanusree88/ViT-MRI-FineTuning/resolve/main/archive%20(5).zip', '')
-    image_paths, labels = prepare_dataset('extracted_folder/')
+    # Define an extraction directory
+    extraction_dir = 'extracted_folder/'
+    os.makedirs(extraction_dir, exist_ok=True)
+
+    # Extract both zip files
+    extract_zip(zip_file_1, extraction_dir)
+    extract_zip(zip_file_2, extraction_dir)
+
+    # Prepare dataset
+    image_paths, labels = prepare_dataset(extraction_dir)
     dataset = CustomImageDataset(image_paths, labels)
     train_loader = DataLoader(dataset, batch_size=32, shuffle=True)
+
+    # Fine-tune the model
     final_loss = fine_tune_model(train_loader)
     st.write(f"Training Complete with Final Loss: {final_loss}")
+