Update app.py

app.py

@@ -95,8 +95,11 @@ class CustomImageDataset(Dataset):
 
 # Training function for classification
 def fine_tune_classification_model(train_loader):
-    model = ResNetForImageClassification.from_pretrained('microsoft/resnet-50', num_labels=3)
+    # Load the ResNet model with ignore_mismatched_sizes
+    model = ResNetForImageClassification.from_pretrained('microsoft/resnet-50', num_labels=3, ignore_mismatched_sizes=True)
+    model.classifier = torch.nn.Linear(model.config.hidden_size, 3)  # Update classifier for 3 labels
     model.train()
+    
     optimizer = AdamW(model.parameters(), lr=1e-4)
     criterion = torch.nn.CrossEntropyLoss()
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -112,6 +115,7 @@ def fine_tune_classification_model(train_loader):
             loss.backward()
             optimizer.step()
             running_loss += loss.item()
+    
     return running_loss / len(train_loader)
 
 # Streamlit UI for Fine-tuning
@@ -140,8 +144,10 @@ if st.button("Start Training"):
 
 # Segmentation function (using SegFormer)
 def fine_tune_segmentation_model(train_loader):
-    model = SegformerForImageSegmentation.from_pretrained('nvidia/segformer-b0', num_labels=3)
+    # Load the Segformer model with ignore_mismatched_sizes
+    model = SegformerForSemanticSegmentation.from_pretrained('nvidia/segformer-b0', num_labels=3, ignore_mismatched_sizes=True)
     model.train()
+    
     optimizer = AdamW(model.parameters(), lr=1e-4)
     criterion = torch.nn.CrossEntropyLoss()
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -157,6 +163,7 @@ def fine_tune_segmentation_model(train_loader):
             loss.backward()
             optimizer.step()
             running_loss += loss.item()
+    
     return running_loss / len(train_loader)
 
 # Add a button for segmentation training