add modified files

src/app.py

@@ -24,7 +24,7 @@ def compute_saliency_map(model, input_tensor, method):
     input_tensor = input_tensor.to(config.DEVICE)
     input_tensor.requires_grad_()
 
-    # Get prediction
+
     output = model(input_tensor)
     pred_class = output.argmax(dim=1).item()
     confidence = torch.softmax(output, dim=1)[0][pred_class].item()
@@ -35,7 +35,12 @@ def compute_saliency_map(model, input_tensor, method):
     elif method == "smoothgrad":
         attr = NoiseTunnel(Saliency(model))
         attributions = attr.attribute(
-            input_tensor, nt_type="smoothgrad", target=pred_class, nt_samples=20, stdevs=0.2)
+            input_tensor, 
+            nt_type="smoothgrad", 
+            target=pred_class, 
+            nt_samples=config.SMOOTHGRAD_SAMPLES, 
+            stdevs=config.SMOOTHGRAD_STDEV
+        )
     elif method == "guided":
         attr = GuidedBackprop(model)
         attributions = attr.attribute(input_tensor, target=pred_class)
@@ -43,7 +48,7 @@ def compute_saliency_map(model, input_tensor, method):
         raise ValueError("Unsupported method")
 
     saliency = attributions.squeeze().abs().cpu().detach().numpy()
-    saliency = np.max(saliency, axis=0)  # to grayscale
+    saliency = np.max(saliency, axis=0)  
 
     return pred_class, confidence, saliency
 
@@ -68,7 +73,7 @@ def run_saliency(model, input_tensor):
 
     output[0, pred_class].backward()
     saliency = input_tensor.grad.abs().squeeze().cpu().numpy()
-    saliency = np.max(saliency, axis=0)  # convert to grayscale
+    saliency = np.max(saliency, axis=0)  
     return pred_class, confidence, saliency
 
 
@@ -77,7 +82,7 @@ def get_saliency_figure(input_tensor, saliency_map):
     saliency_map /= saliency_map.max() + 1e-10
 
     img_np = input_tensor.squeeze().detach().cpu().numpy()
-    img_np = np.transpose(img_np, (1, 2, 0))  # C,H,W → H,W,C
+    img_np = np.transpose(img_np, (1, 2, 0))  
     img_np = (img_np * 0.5 + 0.5).clip(0, 1)
 
     saliency_rgb = np.stack([saliency_map]*3, axis=-1)

src/config.py

@@ -19,3 +19,38 @@ DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"
 
 
 NUM_WORKERS = 2
+
+
+TUNING_EPOCHS = 5
+TUNING_TRIALS = 10
+TUNING_BATCH_SIZE = 32
+
+
+LR_SCHEDULER_PATIENCE = 2
+LR_SCHEDULER_FACTOR = 0.5
+
+
+WEIGHT_DECAY = 1e-4
+DROPOUT_RATE = 0.3
+
+
+DATA_AUG_ROTATION = 15
+DATA_AUG_COLOR_JITTER = 0.1
+DATA_AUG_TRANSLATE = 0.1
+DATA_AUG_SCALE = (0.8, 1.0)
+
+
+GRAD_CLIP_VALUE = 1.0
+
+
+SALIENCY_METHODS = ["saliency", "smoothgrad", "guided"]
+SMOOTHGRAD_SAMPLES = 20
+SMOOTHGRAD_STDEV = 0.2
+
+
+INFERENCE_DIR = os.path.join(DATA_DIR, "inference_test")
+
+
+os.makedirs(LOG_DIR, exist_ok=True)
+os.makedirs(os.path.dirname(MODEL_SAVE_PATH), exist_ok=True)
+os.makedirs(INFERENCE_DIR, exist_ok=True)

src/data_loader.py

@@ -1,13 +1,22 @@
 import os
+import logging
 from torchvision import datasets, transforms
 from torch.utils.data import DataLoader
+from src import config
 
-def get_transforms(image_size=224):
+def get_transforms(image_size=config.IMAGE_SIZE):
     train_transforms = transforms.Compose([
         transforms.Resize((image_size, image_size)),
         transforms.RandomHorizontalFlip(),
-        transforms.RandomRotation(10),
-        transforms.ColorJitter(brightness=0.2, contrast=0.2),
+        transforms.RandomRotation(config.DATA_AUG_ROTATION),
+        transforms.ColorJitter(
+            brightness=config.DATA_AUG_COLOR_JITTER, 
+            contrast=config.DATA_AUG_COLOR_JITTER, 
+            saturation=config.DATA_AUG_COLOR_JITTER, 
+            hue=config.DATA_AUG_COLOR_JITTER
+        ),
+        transforms.RandomAffine(degrees=0, translate=(config.DATA_AUG_TRANSLATE, config.DATA_AUG_TRANSLATE)),
+        transforms.RandomResizedCrop(image_size, scale=config.DATA_AUG_SCALE),
         transforms.ToTensor(),
         transforms.Normalize([0.5]*3, [0.5]*3)
     ])
@@ -20,21 +29,31 @@ def get_transforms(image_size=224):
 
     return train_transforms, val_test_transforms
 
-def get_dataloaders(data_dir, batch_size=32, image_size=224, num_workers=2):
+def get_dataloaders(data_dir, batch_size=config.BATCH_SIZE, image_size=config.IMAGE_SIZE, num_workers=config.NUM_WORKERS):
     train_transforms, val_test_transforms = get_transforms(image_size)
 
     train_dir = os.path.join(data_dir, 'train')
     val_dir = os.path.join(data_dir, 'val')
     test_dir = os.path.join(data_dir, 'test')
 
+    logging.info(f"Loading datasets from: {data_dir}")
+    logging.info(f"Train directory: {train_dir}")
+    logging.info(f"Validation directory: {val_dir}")
+    logging.info(f"Test directory: {test_dir}")
+
     train_dataset = datasets.ImageFolder(train_dir, transform=train_transforms)
     val_dataset = datasets.ImageFolder(val_dir, transform=val_test_transforms)
     test_dataset = datasets.ImageFolder(test_dir, transform=val_test_transforms)
 
+    logging.info(f"Train dataset size: {len(train_dataset)}")
+    logging.info(f"Validation dataset size: {len(val_dataset)}")
+    logging.info(f"Test dataset size: {len(test_dataset)}")
+
     train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
     val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
     test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
 
     class_names = train_dataset.classes
+    logging.info(f"Classes: {class_names}")
 
     return train_loader, val_loader, test_loader, class_names

src/ensemble.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+from src.model import TrashNetClassifier
+from torchvision.models import resnet18, efficientnet_b0
+
+class EnsembleModel(nn.Module):
+    def __init__(self, num_classes=6):
+        super(EnsembleModel, self).__init__()
+        
+
+        self.model1 = TrashNetClassifier(num_classes=num_classes)
+        
+
+        self.model2 = resnet18(pretrained=True)
+        self.model2.fc = nn.Linear(self.model2.fc.in_features, num_classes)
+        
+
+        self.model3 = efficientnet_b0(pretrained=True)
+        self.model3.classifier[1] = nn.Linear(self.model3.classifier[1].in_features, num_classes)
+        
+    def forward(self, x):
+
+        out1 = self.model1(x)
+        out2 = self.model2(x)
+        out3 = self.model3(x)
+        
+
+        return (out1 + out2 + out3) / 3

src/hyperparameter_tuning.py

@@ -0,0 +1,165 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+import random
+from src.model import TrashNetClassifier
+from src.data_loader import get_dataloaders
+from src import config
+
+
+import logging
+import time
+from datetime import datetime
+import os
+
+
+def setup_tuning_logging(log_dir):
+    os.makedirs(log_dir, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    log_file = os.path.join(log_dir, f"hyperparameter_tuning_{timestamp}.log")
+    
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(levelname)s - %(message)s',
+        handlers=[
+            logging.FileHandler(log_file),
+            logging.StreamHandler()
+        ]
+    )
+    return log_file
+
+def train_model_for_validation(model, train_loader, val_loader, lr, weight_decay, device, epochs=config.TUNING_EPOCHS):
+    model = model.to(device)
+    
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(
+        model.parameters(), 
+        lr=lr,
+        weight_decay=weight_decay
+    )
+    
+    best_val_acc = 0.0
+    
+    logging.info(f"Starting validation training with lr={lr}, weight_decay={weight_decay}")
+    
+    for epoch in range(epochs):
+
+        model.train()
+        running_loss, running_acc = 0.0, 0.0
+        for batch_idx, (images, labels) in enumerate(train_loader):
+            if batch_idx % 20 == 0:
+                logging.info(f"  Epoch {epoch+1}/{epochs}, Batch {batch_idx}/{len(train_loader)}")
+                
+            images, labels = images.to(device), labels.to(device)
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            
+            preds = torch.argmax(outputs, dim=1)
+            acc = (preds == labels).float().mean()
+            running_loss += loss.item()
+            running_acc += acc.item()
+        
+        train_loss = running_loss / len(train_loader)
+        train_acc = running_acc / len(train_loader)
+        
+
+        model.eval()
+        val_loss, val_acc = 0.0, 0.0
+        with torch.no_grad():
+            for images, labels in val_loader:
+                images, labels = images.to(device), labels.to(device)
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+                
+                preds = torch.argmax(outputs, dim=1)
+                acc = (preds == labels).float().mean()
+                val_loss += loss.item()
+                val_acc += acc.item()
+        
+        val_loss /= len(val_loader)
+        val_acc /= len(val_loader)
+        
+        logging.info(f"  Epoch {epoch+1}/{epochs}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.4f}, Val Loss={val_loss:.4f}, Val Acc={val_acc:.4f}")
+        
+        if val_acc > best_val_acc:
+            best_val_acc = val_acc
+            logging.info(f"  New best validation accuracy: {best_val_acc:.4f}")
+    
+    return best_val_acc
+
+def run_hyperparameter_search():
+
+    log_file = setup_tuning_logging(config.LOG_DIR)
+    logging.info(f"Hyperparameter tuning logs will be saved to: {log_file}")
+    
+    device = torch.device(config.DEVICE)
+    logging.info(f"Using device: {device}")
+    
+
+    logging.info("Loading datasets...")
+    train_loader, val_loader, _, class_names = get_dataloaders(
+        data_dir=config.DATA_DIR,
+        batch_size=config.TUNING_BATCH_SIZE,
+        image_size=config.IMAGE_SIZE,
+        num_workers=config.NUM_WORKERS
+    )
+    
+
+    learning_rates = [1e-5, 1e-4, 5e-4, 1e-3]
+    weight_decays = [1e-5, 1e-4, 1e-3]
+    
+
+    num_trials = config.TUNING_TRIALS
+    
+    best_acc = 0.0
+    best_config = {"lr": 0, "weight_decay": 0}
+    
+    logging.info("Starting hyperparameter search...")
+    logging.info(f"Number of trials: {num_trials}")
+    logging.info(f"Learning rates to try: {learning_rates}")
+    logging.info(f"Weight decays to try: {weight_decays}")
+    
+    start_time = time.time()
+    
+    for trial in range(num_trials):
+        trial_start = time.time()
+
+        lr = random.choice(learning_rates)
+        weight_decay = random.choice(weight_decays)
+        
+        logging.info(f"\nTrial {trial+1}/{num_trials}")
+        logging.info(f"Testing lr={lr}, weight_decay={weight_decay}")
+        
+
+        model = TrashNetClassifier(num_classes=len(class_names))
+        
+
+        val_acc = train_model_for_validation(
+            model=model,
+            train_loader=train_loader,
+            val_loader=val_loader,
+            lr=lr,
+            weight_decay=weight_decay,
+            device=device
+        )
+        
+        trial_time = time.time() - trial_start
+        logging.info(f"Trial {trial+1} completed in {trial_time:.2f}s")
+        logging.info(f"Validation accuracy: {val_acc:.4f}")
+        
+
+        if val_acc > best_acc:
+            best_acc = val_acc
+            best_config = {"lr": lr, "weight_decay": weight_decay}
+            logging.info(f"New best config found!")
+    
+    total_time = time.time() - start_time
+    logging.info(f"\nHyperparameter search completed in {total_time:.2f}s")
+    logging.info(f"Best config: lr={best_config['lr']}, weight_decay={best_config['weight_decay']}")
+    logging.info(f"Best validation accuracy: {best_acc:.4f}")
+    
+    return best_config

src/model.py

@@ -2,25 +2,29 @@ import torch
 import torch.nn as nn
 import torchvision.models as models
 from src import config
+from torchvision.models import mobilenet_v2
+
 
 class TrashNetClassifier(nn.Module):
-    def __init__(self):
+    def __init__(self, num_classes=config.NUM_CLASSES):
         super(TrashNetClassifier, self).__init__()
+        self.backbone = mobilenet_v2(pretrained=True)
         
-        self.backbone = models.mobilenet_v2(pretrained=True)
-        
+
         if config.FREEZE_BACKBONE:
-            for param in self.backbone.features.parameters():
+            for param in list(self.backbone.parameters())[:-8]:
                 param.requires_grad = False
-                
-        self.backbone.classifier = nn.Sequential(
-            nn.Dropout(p=0.3),
-            nn.Linear(self.backbone.last_channel, 256),
-            nn.ReLU(),
-            nn.BatchNorm1d(256),
-            nn.Dropout(p=0.3),
-            nn.Linear(256, config.NUM_CLASSES)
+            
+
+        in_features = self.backbone.classifier[1].in_features
+        self.backbone.classifier = nn.Identity()
+        
+        self.classifier = nn.Sequential(
+            nn.Dropout(config.DROPOUT_RATE),  
+            nn.Linear(in_features, num_classes)
         )
         
     def forward(self, x):
-        return self.backbone(x)
+        x = self.backbone(x)
+        x = self.classifier(x)
+        return x

src/predict.py

@@ -19,7 +19,7 @@ def preprocess_image(image_path, image_size):
         transforms.Normalize([0.5]*3, [0.5]*3)
     ])
     image = Image.open(image_path).convert("RGB")
-    return transform(image).unsqueeze(0)  # [1, C, H, W]
+    return transform(image).unsqueeze(0)  
 
 def predict_image(model, image_tensor, class_names, device):
     image_tensor = image_tensor.to(device)

src/train.py

@@ -1,3 +1,6 @@
+from datetime import datetime
+import os
+import logging
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -5,16 +8,39 @@ from src import config
 import time
 from torch.utils.tensorboard import SummaryWriter
 
+
 def calculate_accuracy(y_pred, y_true):
     preds = torch.argmax(y_pred, dim=1)
     correct = (preds == y_true).sum().item()
     return correct / len(y_true)
 
+
+def setup_logging(log_dir):
+    os.makedirs(log_dir, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    log_file = os.path.join(log_dir, f"training_{timestamp}.log")
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(levelname)s - %(message)s',
+        handlers=[
+            logging.FileHandler(log_file),
+            logging.StreamHandler()
+        ]
+    )
+    return log_file
+
+
 def train_one_epoch(model, dataloader, criterion, optimizer, device):
     model.train()
     running_loss, running_acc = 0.0, 0.0
+    batch_count = len(dataloader)
+
+    logging.info(f"Training on {batch_count} batches")
+    for batch_idx, (images, labels) in enumerate(dataloader):
+        if batch_idx % 10 == 0:
+            logging.info(f"  Batch {batch_idx}/{batch_count}")
 
-    for images, labels in dataloader:
         images, labels = images.to(device), labels.to(device)
 
         optimizer.zero_grad()
@@ -23,6 +49,10 @@ def train_one_epoch(model, dataloader, criterion, optimizer, device):
         acc = calculate_accuracy(outputs, labels)
 
         loss.backward()
+
+        torch.nn.utils.clip_grad_norm_(
+            model.parameters(), max_norm=config.GRAD_CLIP_VALUE)
+
         optimizer.step()
 
         running_loss += loss.item()
@@ -30,54 +60,91 @@ def train_one_epoch(model, dataloader, criterion, optimizer, device):
 
     return running_loss / len(dataloader), running_acc / len(dataloader)
 
-def validate(model, dataloader, criterion, device):
-    model.eval()
-    val_loss, val_acc = 0.0, 0.0
 
-    with torch.no_grad():
-        for images, labels in dataloader:
-            images, labels = images.to(device), labels.to(device)
-            outputs = model(images)
-            loss = criterion(outputs, labels)
-            acc = calculate_accuracy(outputs, labels)
+def train_model(model, train_loader, val_loader, epochs=config.EPOCHS, lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY, device=config.DEVICE):
 
-            val_loss += loss.item()
-            val_acc += acc
+    log_file = setup_logging(config.LOG_DIR)
+    logging.info(f"Training logs will be saved to: {log_file}")
 
-    return val_loss / len(dataloader), val_acc / len(dataloader)
+    logging.info(f"Training configuration:")
+    logging.info(f"  Epochs: {epochs}")
+    logging.info(f"  Learning rate: {lr}")
+    logging.info(f"  Weight decay: {weight_decay}")
+    logging.info(f"  Device: {device}")
+    logging.info(f"  Batch size: {config.BATCH_SIZE}")
+    logging.info(f"  Image size: {config.IMAGE_SIZE}")
 
-def train_model(model, train_loader, val_loader, epochs=config.EPOCHS, lr=config.LEARNING_RATE, device=config.DEVICE):
     model = model.to(device)
-    optimizer = optim.Adam(model.parameters(), lr=lr)
+    optimizer = optim.Adam(model.parameters(), lr=lr,
+                           weight_decay=weight_decay)
+
+    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer,
+        mode='max',
+        factor=config.LR_SCHEDULER_FACTOR,
+        patience=config.LR_SCHEDULER_PATIENCE,
+        verbose=True
+    )
+
     criterion = nn.CrossEntropyLoss()
     best_val_acc = 0.0
-    
- 
 
     run_name = time.strftime("run_%Y%m%d-%H%M")
     log_dir = f"{config.LOG_DIR}/{run_name}"
     writer = SummaryWriter(log_dir=log_dir)
 
-    print(f"Training on: {device.upper()}\n")
+    logging.info(f"Training on: {device.upper()}\n")
 
     for epoch in range(epochs):
-        train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
+        epoch_start_time = time.time()
+        logging.info(f"Epoch {epoch+1}/{epochs} started")
+
+        train_loss, train_acc = train_one_epoch(
+            model, train_loader, criterion, optimizer, device)
+
+        logging.info("Validating...")
         val_loss, val_acc = validate(model, val_loader, criterion, device)
 
-        print(f"Epoch {epoch+1}/{epochs}")
-        print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc*100:.2f}%")
-        print(f"Val   Loss: {val_loss:.4f} | Val   Acc: {val_acc*100:.2f}%")
-        
+        epoch_time = time.time() - epoch_start_time
+
+        scheduler.step(val_acc)
+
+        logging.info(
+            f"Epoch {epoch+1}/{epochs} completed in {epoch_time:.2f}s")
+        logging.info(
+            f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc*100:.2f}%")
+        logging.info(
+            f"Val   Loss: {val_loss:.4f} | Val   Acc: {val_acc*100:.2f}%")
+
         writer.add_scalar("Loss/train", train_loss, epoch)
         writer.add_scalar("Loss/val", val_loss, epoch)
         writer.add_scalar("Accuracy/train", train_acc, epoch)
         writer.add_scalar("Accuracy/val", val_acc, epoch)
 
-
         if val_acc > best_val_acc:
             best_val_acc = val_acc
             torch.save(model.state_dict(), config.MODEL_SAVE_PATH)
-            print("Model saved!\n")
+            logging.info("Model saved!")
 
     writer.close()
-    print("Training complete. Best Val Acc: {:.2f}%".format(best_val_acc * 100))
+    logging.info("Training complete. Best Val Acc: {:.2f}%".format(
+        best_val_acc * 100))
+
+    return best_val_acc
+
+
+def validate(model, dataloader, criterion, device):
+    model.eval()
+    val_loss, val_acc = 0.0, 0.0
+
+    with torch.no_grad():
+        for images, labels in dataloader:
+            images, labels = images.to(device), labels.to(device)
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            acc = calculate_accuracy(outputs, labels)
+
+            val_loss += loss.item()
+            val_acc += acc
+
+    return val_loss / len(dataloader), val_acc / len(dataloader)

src/train_with_tuning.py

@@ -0,0 +1,38 @@
+import torch
+from src.hyperparameter_tuning import run_hyperparameter_search
+from src.model import TrashNetClassifier
+from src.data_loader import get_dataloaders
+from src.train import train_model
+from src import config
+
+if __name__ == "__main__":
+
+    print("Starting hyperparameter search...")
+    best_config = run_hyperparameter_search()
+    
+
+    print("\nTraining with best hyperparameters...")
+    
+
+    train_loader, val_loader, test_loader, class_names = get_dataloaders(
+        data_dir=config.DATA_DIR,
+        batch_size=config.BATCH_SIZE,
+        image_size=config.IMAGE_SIZE,
+        num_workers=config.NUM_WORKERS
+    )
+    
+
+    model = TrashNetClassifier(num_classes=len(class_names))
+    
+
+    train_model(
+        model=model,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        epochs=config.EPOCHS,
+        lr=best_config["lr"],
+        weight_decay=best_config["weight_decay"],
+        device=config.DEVICE
+    )
+    
+    print("Training complete!")