diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 00000000..53276af2 Binary files /dev/null and b/.DS_Store differ diff --git a/Image Classifier Project.html b/Image Classifier Project.html new file mode 100644 index 00000000..f71aa30e --- /dev/null +++ b/Image Classifier Project.html @@ -0,0 +1,8632 @@ + + + + + +Image Classifier Project + + + + + + + + + + + + +
+
+ +
+ + +
+ +
+
+ +
+ +
+ +
+
+ +
+ +
+
+ +
+ + +
+
+ +
+ + +
+
+ +
+ + +
+
+ +
+ +
+ +
+ + +
+
+ +
+ +
+ +
+ + +
+
+ +
+ +
+ + +
+
+ +
+ +
+ + +
+
+ +
+ +
+ + +
+ + +
+ +
+
+ + diff --git a/predict.py b/predict.py index e69de29b..fddb6f5c 100644 --- a/predict.py +++ b/predict.py @@ -0,0 +1,48 @@ +import torch +import argparse +from torchvision import models +from PIL import Image +import numpy as np + +def process_image(image_path): + image = Image.open(image_path) + image = image.resize((256, 256)).crop((16, 16, 240, 240)) + np_image = np.array(image) / 255.0 + mean = np.array([0.485, 0.456, 0.406]) + std = np.array([0.229, 0.224, 0.225]) + np_image = (np_image - mean) / std + np_image = np_image.transpose((2, 0, 1)) + return torch.tensor(np_image).float().unsqueeze(0) + +def predict(args): + device = torch.device('cuda' if args.gpu and torch.cuda.is_available() else 'cpu') + checkpoint = torch.load(args.checkpoint) + model = models.vgg16(pretrained=True) if checkpoint['arch'] == 'vgg16' else models.resnet50(pretrained=True) + model.classifier = checkpoint['classifier'] + model.load_state_dict(checkpoint['model_state_dict']) + model.class_to_idx = checkpoint['class_to_idx'] + model.to(device) + model.eval() + + image = process_image(args.image_path).to(device) + with torch.no_grad(): + output = model(image) + probs, indices = torch.exp(output).topk(args.top_k) + idx_to_class = {v: k for k, v in model.class_to_idx.items()} + classes = [idx_to_class[idx.item()] for idx in indices[0]] + + print("Predictions:") + for prob, cls in zip(probs[0], classes): + print(f"{cls}: {prob.item():.4f}") + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Predict the class of an image') + parser.add_argument('image_path', type=str, help='Path to image') + parser.add_argument('checkpoint', type=str, help='Path to model checkpoint') + parser.add_argument('--top_k', type=int, default=5, help='Return top K classes') + parser.add_argument('--gpu', action='store_true', help='Use GPU for inference') + args = parser.parse_args() + predict(args) + + +# python predict.py flowers/test/1/image_06743.jpg checkpoint.pth --top_k 5 --gpu \ No newline at end of file diff --git a/train.py b/train.py index e69de29b..00c6fcdc 100644 --- a/train.py +++ b/train.py @@ -0,0 +1,82 @@ +# train.py - Script to train a new network on a dataset and save the model as a checkpoint +import torch +import argparse +from torchvision import datasets, transforms, models +from torch import nn, optim +import json +import os + +def get_input_args(): + parser = argparse.ArgumentParser(description='Train a neural network') + parser.add_argument('data_dir', type=str, help='Path to dataset') + parser.add_argument('--save_dir', type=str, default='.', help='Directory to save checkpoints') + parser.add_argument('--arch', type=str, default='vgg16', choices=['vgg16', 'resnet50'], help='Pretrained model architecture') + parser.add_argument('--learning_rate', type=float, default=0.001, help='Learning rate') + parser.add_argument('--hidden_units', type=int, default=512, help='Hidden layer units') + parser.add_argument('--epochs', type=int, default=5, help='Number of epochs') + parser.add_argument('--gpu', action='store_true', help='Use GPU for training') + return parser.parse_args() + +def train_model(args): + device = torch.device('cuda' if args.gpu and torch.cuda.is_available() else 'cpu') + + # Data transformations + data_transforms = transforms.Compose([ + transforms.RandomResizedCrop(224), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + dataset = datasets.ImageFolder(os.path.join(args.data_dir, 'train'), transform=data_transforms) + dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True) + + # Load model + model = models.vgg16(pretrained=True) if args.arch == 'vgg16' else models.resnet50(pretrained=True) + for param in model.parameters(): + param.requires_grad = False + + # Define classifier + classifier = nn.Sequential( + nn.Linear(25088, args.hidden_units), + nn.ReLU(), + nn.Dropout(0.2), + nn.Linear(args.hidden_units, 102), + nn.LogSoftmax(dim=1) + ) + model.classifier = classifier + model.to(device) + + criterion = nn.NLLLoss() + optimizer = optim.Adam(model.classifier.parameters(), lr=args.learning_rate) + + # Training loop + for epoch in range(args.epochs): + running_loss = 0 + for inputs, labels in dataloader: + inputs, labels = inputs.to(device), labels.to(device) + optimizer.zero_grad() + outputs = model(inputs) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + running_loss += loss.item() + print(f"Epoch {epoch+1}/{args.epochs} - Loss: {running_loss/len(dataloader):.4f}") + + # Save checkpoint + model.class_to_idx = dataset.class_to_idx + checkpoint = { + 'arch': args.arch, + 'model_state_dict': model.state_dict(), + 'class_to_idx': model.class_to_idx, + 'classifier': model.classifier + } + torch.save(checkpoint, os.path.join(args.save_dir, 'checkpoint.pth')) + print("Checkpoint saved!") + +if __name__ == '__main__': + args = get_input_args() + train_model(args) + + + +# python train.py flowers --save_dir ./ --arch vgg16 --learning_rate 0.001 --hidden_units 512 --epochs 5 --gpu \ No newline at end of file