Update app.py

app.py

@@ -54,154 +54,7 @@ else:
     print(f"Dataset zip file '{dataset_zip}' not found.")
 
 
-# Path to the data directory
-data_dir = '/home/user/app/data'
 
-# Define data transformations
-data_transforms = {
-    'train': transforms.Compose([
-        transforms.RandomResizedCrop(224),
-        transforms.RandomRotation(15),
-        transforms.RandomHorizontalFlip(),
-        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ]),
-    'valid': transforms.Compose([
-        transforms.Resize(256),
-        transforms.CenterCrop(224),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ]),
-}
-
-# Create the datasets from the image folder
-image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
-                  for x in ['train', 'valid']}
-
-# Create the dataloaders
-dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=32, shuffle=True, num_workers=4)
-               for x in ['train', 'valid']}
-
-# Class names
-class_names = image_datasets['train'].classes
-print(f"Classes: {class_names}")
-
-# Check if a GPU is available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Load pre-trained ResNet50 model
-model = models.resnet50(weights='ResNet50_Weights.DEFAULT')  # Use weights instead of pretrained
-
-# Modify the final layer to match the number of classes
-num_ftrs = model.fc.in_features
-model.fc = nn.Linear(num_ftrs, len(class_names))  # Output classes match
-
-# Move the model to the GPU if available
-model = model.to(device)
-
-# Loss function and optimizer
-criterion = nn.CrossEntropyLoss()
-optimizer = optim.Adam(model.parameters(), lr=0.001)
-
-# Learning rate scheduler
-scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
-
-# Number of epochs
-num_epochs = 10
-
-# Training function with detailed output for each epoch
-def train_model(model, criterion, optimizer, scheduler, num_epochs=10):
-    since = time.time()
-    
-    best_model_wts = model.state_dict()
-    best_acc = 0.0
-
-    for epoch in range(num_epochs):
-        epoch_start = time.time()  # Start time for this epoch
-        print(f'Epoch {epoch + 1}/{num_epochs}')
-        print('-' * 10)
-
-        # Each epoch has a training and validation phase
-        for phase in ['train', 'valid']:
-            if phase == 'train':
-                model.train()  # Set model to training mode
-            else:
-                model.eval()   # Set model to evaluate mode
-
-            running_loss = 0.0
-            running_corrects = 0
-
-            # Iterate over data
-            for inputs, labels in dataloaders[phase]:
-                inputs = inputs.to(device)
-                labels = labels.to(device)
-
-                # Zero the parameter gradients
-                optimizer.zero_grad()
-
-                # Forward
-                with torch.set_grad_enabled(phase == 'train'):
-                    outputs = model(inputs)
-                    _, preds = torch.max(outputs, 1)
-                    loss = criterion(outputs, labels)
-
-                    # Backward + optimize only if in training phase
-                    if phase == 'train':
-                        loss.backward()
-                        optimizer.step()
-
-                # Statistics
-                running_loss += loss.item() * inputs.size(0)
-                running_corrects += torch.sum(preds == labels.data)
-
-            if phase == 'train':
-                scheduler.step()
-
-            # Calculate epoch loss and accuracy
-            epoch_loss = running_loss / len(image_datasets[phase])
-            epoch_acc = running_corrects.double() / len(image_datasets[phase])
-
-            # Print loss and accuracy for each phase
-            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
-
-            # Deep copy the model if it's the best accuracy
-            if phase == 'valid' and epoch_acc > best_acc:
-                best_acc = epoch_acc
-                best_model_wts = model.state_dict()
-
-        epoch_end = time.time()  # End time for this epoch
-        print(f'Epoch {epoch + 1} completed in {epoch_end - epoch_start:.2f} seconds.')
-
-    time_elapsed = time.time() - since
-    print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
-    print(f'Best val Acc: {best_acc:.4f}')
-
-    # Load best model weights
-    model.load_state_dict(best_model_wts)
-    return model
-
-# Train the model
-best_model = train_model(model, criterion, optimizer, scheduler, num_epochs=num_epochs)
-
-# Save the model
-torch.save(model.state_dict(), 'resnet50_garbage_classification.pth')
-
-import pickle
-
-# Manually creating the history dictionary based on the logs you provided
-history = {
-    'train_loss': [1.0083, 0.7347, 0.6510, 0.5762, 0.5478, 0.5223, 0.4974, 0.3464, 0.2896, 0.2604],
-    'train_acc': [0.6850, 0.7687, 0.7913, 0.8126, 0.8210, 0.8272, 0.8355, 0.8870, 0.9049, 0.9136],
-    'val_loss': [0.6304, 0.8616, 0.5594, 0.4006, 0.3968, 0.4051, 0.3223, 0.2221, 0.2125, 0.2076],
-    'val_acc': [0.7985, 0.7307, 0.8260, 0.8655, 0.8793, 0.8729, 0.9094, 0.9338, 0.9338, 0.9326]
-}
-
-# Save the history as a pickle file
-with open('training_history.pkl', 'wb') as f:
-    pickle.dump(history, f)
-
-print('Training history saved as training_history.pkl')