| |
| |
| |
| |
|
|
| import argparse, os |
| from pathlib import Path |
| from PIL import Image |
| import torch, torch.nn as nn, torch.optim as optim |
| from torchvision import transforms, models |
| from torchvision.datasets import ImageFolder |
| from torch.utils.data import DataLoader |
|
|
| def get_loaders(train_dir, val_dir, batch=16, img=224): |
| tr = transforms.Compose([transforms.RandomResizedCrop(img), transforms.RandomHorizontalFlip(), |
| transforms.ToTensor(), transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])]) |
| val = transforms.Compose([transforms.Resize(int(img*1.14)), transforms.CenterCrop(img), |
| transforms.ToTensor(), transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])]) |
| train_ds = ImageFolder(train_dir, transform=tr) |
| val_ds = ImageFolder(val_dir, transform=val) |
| return DataLoader(train_ds, batch_size=batch, shuffle=True), DataLoader(val_ds, batch_size=batch), train_ds.classes |
|
|
| def train(args): |
| train_loader, val_loader, classes = get_loaders(args.train_dir, args.val_dir, args.batch) |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model = models.resnet18(pretrained=True) |
| model.fc = nn.Linear(model.fc.in_features, len(classes)) |
| model.to(device) |
| opt = optim.Adam(model.parameters(), lr=1e-4) |
| loss_fn = nn.CrossEntropyLoss() |
| best=0 |
| for e in range(args.epochs): |
| model.train() |
| for imgs, lbl in train_loader: |
| imgs, lbl = imgs.to(device), lbl.to(device) |
| opt.zero_grad() |
| out = model(imgs) |
| loss = loss_fn(out, lbl) |
| loss.backward(); opt.step() |
| |
| model.eval() |
| correct=total=0 |
| with torch.no_grad(): |
| for imgs,lbl in val_loader: |
| imgs,lbl = imgs.to(device), lbl.to(device) |
| out = model(imgs).argmax(1) |
| correct += (out==lbl).sum().item(); total += lbl.size(0) |
| acc = correct/total if total else 0 |
| print(f"Epoch {e+1}/{args.epochs} val_acc={acc:.4f}") |
| if acc>best: |
| best=acc |
| torch.save({'model':model.state_dict(),'classes':classes}, args.ckpt) |
| print("Best val acc", best) |
|
|
| def predict(args): |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| ck = torch.load(args.ckpt, map_location=device) |
| classes = ck['classes'] |
| model = models.resnet18(pretrained=False) |
| model.fc = nn.Linear(model.fc.in_features, len(classes)) |
| model.load_state_dict(ck['model']); model.to(device).eval() |
| tf = transforms.Compose([transforms.Resize(256), transforms.CenterCrop(224), |
| transforms.ToTensor(), transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])]) |
| img = Image.open(args.image).convert('RGB') |
| x = tf(img).unsqueeze(0).to(device) |
| with torch.no_grad(): |
| p = model(x).argmax(1).item() |
| print("Predicted:", classes[p]) |
|
|
| if __name__=="__main__": |
| p=argparse.ArgumentParser() |
| p.add_argument('--mode', choices=['train','predict'], required=True) |
| p.add_argument('--train_dir', default='data/train') |
| p.add_argument('--val_dir', default='data/val') |
| p.add_argument('--epochs', type=int, default=3) |
| p.add_argument('--batch', type=int, default=16) |
| p.add_argument('--ckpt', default='ckpt.pth') |
| p.add_argument('--image') |
| args=p.parse_args() |
| if args.mode=='train': |
| train(args) |
| else: |
| if not args.image: raise SystemExit("Provide --image for predict") |
| predict(args) |
|
|
