app.py

import os
import shutil
import subprocess
import zipfile
import time
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.optim import lr_scheduler
import subprocess
import zipfile
from PIL import Image
import gradio as gr

# Step 1: Setup Kaggle API
# Ensure the .kaggle directory exists
kaggle_dir = os.path.expanduser("~/.kaggle")
if not os.path.exists(kaggle_dir):
    os.makedirs(kaggle_dir)

# Step 2: Copy the kaggle.json file to the ~/.kaggle directory
kaggle_json_path = "kaggle.json"
kaggle_dest_path = os.path.join(kaggle_dir, "kaggle.json")

if not os.path.exists(kaggle_dest_path):
    shutil.copy(kaggle_json_path, kaggle_dest_path)
    os.chmod(kaggle_dest_path, 0o600)
    print("Kaggle API key copied and permissions set.")
else:
    print("Kaggle API key already exists.")
    
# Step 3: Download the dataset from Kaggle using Kaggle CLI
dataset_name = "mostafaabla/garbage-classification"
print(f"Downloading the dataset: {dataset_name}")
download_command = f"kaggle datasets download -d {dataset_name}"

# Run the download command
subprocess.run(download_command, shell=True)

# Step 4: Unzip the downloaded dataset
dataset_zip = "garbage-classification.zip"
extracted_folder = "./garbage-classification"

# Check if the zip file exists
if os.path.exists(dataset_zip):
    if not os.path.exists(extracted_folder):
        with zipfile.ZipFile(dataset_zip, 'r') as zip_ref:
            zip_ref.extractall(extracted_folder)
            print("Dataset unzipped successfully!")
    else:
        print("Dataset already unzipped.")
else:
    print(f"Dataset zip file '{dataset_zip}' not found.")


# Path to the data directory
data_dir = 'C:\\Users\\kendr\\Downloads\\data'  # Adjust this if necessary

# 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')



# Load your model
def load_model():
    model = models.resnet50(weights='DEFAULT')  # Using default weights for initialization
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 12)  # Adjust to the number of classes you have
    
    # Load the state dict without the weights_only argument
    model.load_state_dict(torch.load('resnet50_garbage_classification.pth', map_location=torch.device('cpu')))
    
    model.eval()  # Set to evaluation mode
    return model

model = load_model()

# Define image transformations
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# Class names
class_names = ['battery', 'biological', 'brown-glass', 'cardboard', 
               'clothes', 'green-glass', 'metal', 'paper', 
               'plastic', 'shoes', 'trash', 'white-glass']

# Define bin colors for each class
bin_colors = {
    'battery': 'Merah (Red)',        # Limbah berbahaya
    'biological': 'Cokelat (Brown)',  # Limbah organik
    'brown-glass': 'Hijau (Green)',   # Gelas berwarna coklat
    'cardboard': 'Kuning (Yellow)',    # Limbah daur ulang
    'clothes': 'Biru (Blue)',         # Pakaian dan tekstil
    'green-glass': 'Hijau (Green)',    # Gelas berwarna hijau
    'metal': 'Kuning (Yellow)',        # Limbah daur ulang
    'paper': 'Kuning (Yellow)',        # Limbah daur ulang
    'plastic': 'Kuning (Yellow)',      # Limbah daur ulang
    'shoes': 'Biru (Blue)',           # Pakaian dan tekstil
    'trash': 'Hitam (Black)',         # Limbah umum
    'white-glass': 'Putih (White)'    # Gelas berwarna putih
}

# Define the prediction function
def predict(image):
    image = Image.fromarray(image)  # Convert numpy array to PIL Image
    image = transform(image)  # Apply transformations
    image = image.unsqueeze(0)  # Add batch dimension
    
    with torch.no_grad():
        outputs = model(image)
        _, predicted = torch.max(outputs, 1)
    
    class_name = class_names[predicted.item()]  # Return predicted class name
    bin_color = bin_colors[class_name]  # Get the corresponding bin color
    return class_name, bin_color  # Return both class name and bin color

# Make Gradio Interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Unggah Gambar"),
    outputs=[
        gr.Textbox(label="Jenis Sampah"), 
        gr.Textbox(label="Tong Sampah yang Sesuai")  # 2 output with label
    ],
    title="Klasifikasi Sampah dengan ResNet50",
    description="Unggah gambar sampah, dan model akan mengklasifikasikannya ke dalam salah satu dari 12 kategori bersama dengan warna tempat sampah yang sesuai."
)


iface.launch(share=True)