Update app.py

app.py

@@ -54,181 +54,8 @@ 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 = 20
-
-# 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
-
-# Mengupdate hasil train dan validate terbaru
-history = {
-    'train_loss': [
-        0.9568, 0.6937, 0.5917, 0.5718, 0.5109,
-        0.4824, 0.4697, 0.3318, 0.2785, 0.2680,
-        0.2371, 0.2333, 0.2198, 0.2060, 0.1962,
-        0.1951, 0.1880, 0.1912, 0.1811, 0.1810
-    ],
-    'train_acc': [
-        0.7011, 0.7774, 0.8094, 0.8146, 0.8331,
-        0.8452, 0.8447, 0.8899, 0.9068, 0.9114,
-        0.9216, 0.9203, 0.9254, 0.9306, 0.9352,
-        0.9346, 0.9368, 0.9353, 0.9396, 0.9409
-    ],
-    'val_loss': [
-        0.4934, 0.3939, 0.4377, 0.3412, 0.2614,
-        0.2966, 0.2439, 0.1065, 0.0926, 0.0797,
-        0.0738, 0.0639, 0.0555, 0.0560, 0.0490,
-        0.0479, 0.0455, 0.0454, 0.0438, 0.0427
-    ],
-    'val_acc': [
-        0.8481, 0.8734, 0.8663, 0.8915, 0.9172,
-        0.9011, 0.9221, 0.9649, 0.9714, 0.9759,
-        0.9762, 0.9791, 0.9827, 0.9812, 0.9843,
-        0.9850, 0.9852, 0.9854, 0.9854, 0.9866
-    ]
-}
-
-# Simpan history sebagai file pickle
-with open('training_history.pkl', 'wb') as f:
-    pickle.dump(history, f)
-
-print('Training history saved as training_history.pkl')
-
-
-
-
-import torch
-import torch.nn as nn
-from torchvision import models, transforms
 from PIL import Image
 import gradio as gr
 
@@ -305,4 +132,3 @@ iface = gr.Interface(
 
 iface.launch(share=True)  
 
-token = os.getenv("HUGGINGFACE_TOKEN")