# AI Computer Vision Inspection
**2020 LiteOn Award**
**Detect the defect rate of pcba board**
## ResNet Prefromence

## VGG16 Version

# -*- coding: utf-8 -*-
"""liteonvgg30.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1VOgdwXEpfMTGSdw0Ep4lE2fXu_qQKGog
"""

from google.colab import drive
drive.mount('/content/drive')

from __future__ import print_function, division

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.autograd import Variable
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy
from sklearn.metrics import classification_report

plt.ion()

use_gpu = torch.cuda.is_available()
if use_gpu:
    print("Using CUDA")
else:
    print("Error")

data_dir = 'drive/MyDrive/LiteonRacingData'
TRAIN = 'A'
TEST = 'B'

IMAGENET_MEAN = [0.4914, 0.4822, 0.4465]
IMAGENET_STD = [0.2023, 0.1994, 0.201]
data_transforms = {
    TRAIN: transforms.Compose([

        transforms.RandomHorizontalFlip(),
        transforms.RandomVerticalFlip(),
        transforms.RandomPosterize(bits=2),
        transforms.ColorJitter(brightness=1, contrast=1, saturation=1.0, hue=0.5),
        transforms.GaussianBlur(7,3),
        transforms.RandomAffine(degrees=0,translate=(0.1,0.1),fillcolor=None),
        transforms.ToTensor(),
        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
    ]),

    TEST: transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
    ])
}

image_datasets = {
    x: datasets.ImageFolder(
        os.path.join(data_dir, x),
        transform=data_transforms[x]
    )
    for x in [TRAIN, TEST]
}

dataloaders = {
    x: torch.utils.data.DataLoader(
        image_datasets[x], batch_size=8,
        shuffle=True, num_workers=4
    )
    for x in [TRAIN, TEST]
}

dataset_sizes = {x: len(image_datasets[x]) for x in [TRAIN, TEST]}

for x in [TRAIN, TEST]:
    print("Loaded {} images under {}".format(dataset_sizes[x], x))

print("Classes: ")
class_names = image_datasets[TRAIN].classes
print(image_datasets[TRAIN].classes)

# Get a batch of training data
inputs, classes = next(iter(dataloaders[TRAIN]))

from torchvision import models
# Load the pretrained model from pytorch
vgg16 = models.vgg16_bn(pretrained=True)
# Freeze training for all layers
for param in vgg16.features.parameters():
    param.require_grad = False


# Newly created modules have require_grad=True by default
num_features = vgg16.classifier[6].in_features
features = list(vgg16.classifier.children())[:-1] # Remove last layer
features.extend([
                 nn.Linear(num_features, 8)
                 ]) # Add our layer with 4 outputs
vgg16.classifier = nn.Sequential(*features) # Replace the model classifier


vgg16.load_state_dict(torch.load("drive/MyDrive/969282.pt"))
print("Loaded!")

if use_gpu:
    vgg16.cuda() #.cuda() will move everything to the GPU side
weight_CE = torch.FloatTensor([0.7,0.8,1,3,0.8,0.8,2,2])
criterion = nn.CrossEntropyLoss(ignore_index=255,weight=weight_CE.cuda())
optimizer_ft = optim.Adam(vgg16.parameters(), lr=5e-5, betas=[0.5, 0.999])
exp_lr_scheduler = lr_scheduler.CosineAnnealingLR(optimizer_ft, 30, 0)
#optimizer_ft = optim.SGD(vgg16.parameters(), lr=5e-6, momentum=0.9,weight_decay=5e-4)
#exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.95)

def train_model(vgg, criterion, optimizer, scheduler, num_epochs=50):
    best_model_wts = copy.deepcopy(vgg.state_dict())
    best_acc = 0.0

    avg_loss = 0
    avg_acc = 0
    avg_loss_val = 0
    avg_acc_val = 0

    train_batches = len(dataloaders[TRAIN])
    val_batches = len(dataloaders[TEST])

    for epoch in range(num_epochs):
        if epoch==0:
          with open('version3.txt', 'w') as f:
            f.write("epoch:")
            f.write(str(epoch))
        else:
           with open('version3.txt', 'a') as f:
            f.write("epoch:")
            f.write(str(epoch))


        print("Epoch {}/{}".format(epoch, num_epochs))
        print('-' * 10)

        loss_train = 0
        loss_val = 0
        acc_train = 0
        acc_val = 0

        vgg.train(True)

        for i, data in enumerate(dataloaders[TRAIN]):
            if i % 100 == 0:
                print("\rTraining batch {}/{}".format(i, train_batches/5), end='', flush=True)

            # Use half training dataset
            if i >= train_batches /5:
                break

            inputs, labels = data
            #inputs = inputs + torch.rand(inputs.shape)
            if use_gpu:
                inputs, labels = Variable(inputs.cuda()), Variable(labels.cuda())
            else:
                inputs, labels = Variable(inputs), Variable(labels)


            optimizer.zero_grad()

            outputs = vgg(inputs)
            #_______________________
            # m = nn.LogSoftmax()
            # closs = nn.NLLLoss()
            # outputs=m(outputs)
            # loss = closs(outputs, labels)
            #print('output:',output)
            #________________________
            _, preds = torch.max(outputs.data, 1)
            loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step()

            loss_train += loss.item()
            acc_train += torch.sum(preds == labels.data)

            del inputs, labels, outputs, preds
            torch.cuda.empty_cache()

        print()
        # * 2 as we only used half of the dataset
        avg_loss = loss_train * 8 / dataset_sizes[TRAIN]
        avg_acc = acc_train * 8 / dataset_sizes[TRAIN]

        vgg.train(False)
        vgg.eval()
        #_____
        y_pred = []
        y_true = []
        target_names = ['class 0', 'class 1', 'class 2', 'class 3', 'class 4', 'class 5', 'class 6', 'class 7']
        #_____
        for i, data in enumerate(dataloaders[TEST]):
            if i % 100 == 0:
                print("\rValidation batch {}/{}".format(i, val_batches), end='', flush=True)

            inputs, labels = data

            if use_gpu:
                inputs, labels = Variable(inputs.cuda(), volatile=True), Variable(labels.cuda(), volatile=True)
            else:
                inputs, labels = Variable(inputs, volatile=True), Variable(labels, volatile=True)

            optimizer.zero_grad()

            outputs = vgg(inputs)

            #_______________________
            # m = nn.LogSoftmax()
            # closs = nn.NLLLoss()
            # outputs=m(outputs)
            # loss = closs(outputs, labels)
            #print('output:',output)
            #________________________
            _, preds = torch.max(outputs.data, 1)
            loss = criterion(outputs, labels)

            loss_val += loss.item()
            acc_val += torch.sum(preds == labels.data)
                #____
            preds = preds.cpu().numpy()
            labels = labels.cpu().numpy()
            for j in range(0, len(labels)):
              # print(j)
              y_pred.append(preds[j])
              y_true.append(labels[j])
            # print(y_pred)
            #____
            # del inputs, labels, outputs, preds
            # torch.cuda.empty_cache()
        #print(classification_report(y_true, y_pred, target_names=target_names))

        with open('version3.txt', 'a') as f:
          f.write(classification_report(y_true, y_pred, target_names=target_names))
          f.write('_____________________________________________________________\n')
        #files.download('version1.txt')

        avg_loss_val = loss_val / dataset_sizes[TEST]
        avg_acc_val = acc_val / dataset_sizes[TEST]

        print()
        print("Epoch {} result: ".format(epoch))
        print("Avg loss (train): {:.4f}".format(avg_loss))
        print("Avg acc (train): {:.4f}".format(avg_acc))
        print("Avg loss (val): {:.4f}".format(avg_loss_val))
        print("Avg acc (val): {:.4f}".format(avg_acc_val))
        print('-' * 10)
        print()
        #eval_model(vgg16, criterion)


        if avg_acc_val >= best_acc:
            best_acc = avg_acc_val
            best_model_wts = copy.deepcopy(vgg.state_dict())
            torch.save(vgg16.state_dict(), 'VGG16_18('+str(epoch)+').pt')

    print()
    print("Best acc: {:.4f}".format(best_acc))

    vgg.load_state_dict(best_model_wts)
    return vgg

vgg16 = train_model(vgg16, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=50)
torch.save(vgg16.state_dict(), 'VGG16_3.pt')