# -*- coding: utf-8 -*-
"""「「「「ML_FinalProject.ipynb」的副本」的副本」的副本」的副本

Automatically generated by Colaboratory.

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

# **Download the dataset**
首先先下載需要用到的資料，解壓縮後資料夾內的結構如下
```
finalProject/
|----train/train
    |----7b1c3f70760d2fc744d6fcebfa34d374.jpg
    |----7b1fdeeff00ce7b316a9d61b435ab7fe.jpg
    ...
|----unlabeled_data
    |----AafJpnHwok6965511956.jpg
    |----AaFWonxFbM4044591286.jpg
    ...
|----test/test
    |----000d70d25191ad64f00ca88a227c5985.jpg
    |----00a114110ac6f96c5ac10e792db5197f.jpg
    ...
|----train_data.csv
|----val_data.csv
```
Training data總共有9297筆   
Testing data總共有2800筆
"""

# from google.colab import drive
# drive._mount('/content/drive')

!gdown --id 1U5KGRFT8gTix3eL1hvCxTojcNRuoN5RQ --output "data.zip"
!unzip -q "data.zip"

"""# **Import Packages**"""

import os
import glob
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torch.utils.data as data
import torchvision.transforms as transforms
from torchvision import datasets
from torchvision.datasets import DatasetFolder
from torchvision.transforms import AutoAugmentPolicy
from torch import optim
from torch.utils.data import Dataset, DataLoader, ConcatDataset, Subset
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
import random
import tensorflow as tf
import csv
import pandas as pd

"""### 固定 random seed"""

def same_seeds(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
same_seeds(87)

"""# **Dataset, Data Loader, and Transforms**
Reference:  
https://blog.csdn.net/weixin_41469023/article/details/115375927  
https://www.itread01.com/content/1544541602.html 
"""

train = pd.read_csv('train_data.csv')
train.head()

train_file = train["Name"]
train_file = [os.path.join("train/train/",i) for i in train_file]

train_label = train["Type"]
label_train = []
for i in range(len(train_label)):
    label_train.append(train_label[i])
print(label_train)

class IMAGE(Dataset):
    def __init__(self, root, transform=None):
        self.transform = transform
        if root == 'train/train/':
            self.filenames = train_file
            self.target = label_train
        if root == 'unlabeled_data/':
            self.filenames = glob.glob(os.path.join(root,'*'))
            self.target = [0 for i in range(len(self.filenames))]     
        if root == 'test/test/':
            self.filenames = glob.glob(os.path.join(root,'*'))
            self.target = [0 for i in range(len(self.filenames))]
        self.len = len(self.filenames)
    def __getitem__(self, index):
        image_fn = self.filenames[index]
        image = Image.open(image_fn)
        if image.mode != 'RGB':
            image = image.convert('RGB')
        label = self.target[index]
        if self.transform is not None:
            image = self.transform(image)
        image = image.resize_([3, 256, 256])
        return image, label
    def __len__(self):
        return self.len

"""## Data Augmentation
對現有圖片進行平移、翻轉、旋轉、改變明度對比度等方式來增加資料集
"""

randomAugment = [transforms.ColorJitter(0.1, 0.1, 0.1), transforms.ColorJitter(0.2, 0.2, 0.2)]
autoAugment = [transforms.AutoAugment(policy=AutoAugmentPolicy.IMAGENET),
         transforms.AutoAugment(policy=AutoAugmentPolicy.CIFAR10),
         transforms.AutoAugment(policy=AutoAugmentPolicy.SVHN)]

train_tfm1 = transforms.Compose([
	transforms.Resize((256, 256)),
  # transforms.Resize((128, 128)),
  transforms.RandomChoice(randomAugment),
  transforms.RandomHorizontalFlip(p=0.7),
	transforms.ToTensor(),
  transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_tfm2 = transforms.Compose([
	transforms.Resize((256, 256)),
  # transforms.Resize((128, 128)),
  transforms.RandomChoice(randomAugment),
  transforms.RandomHorizontalFlip(p=0.5),
  transforms.AutoAugment(policy=AutoAugmentPolicy.IMAGENET),
  transforms.RandomAffine(degrees=30, translate=(0.12, 0.12), scale=(0.8, 1.2)),
  # transforms.RandomApply(transform_set, p=0.7),
	transforms.ToTensor(),
  transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_tfm3 = transforms.Compose([
	transforms.Resize((256, 256)),
  # transforms.Resize((128, 128)),
  transforms.RandomChoice(randomAugment),
  transforms.RandomHorizontalFlip(p=0.7),
  transforms.AutoAugment(policy=AutoAugmentPolicy.CIFAR10),
  transforms.RandomAffine(degrees=45, translate=(0.12, 0.12), scale=(0.8, 1.2)),
	transforms.ToTensor(),
  transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_tfm4 = transforms.Compose([
	transforms.Resize((256, 256)),
  # transforms.Resize((128, 128)),
  transforms.RandomChoice(randomAugment),
  transforms.RandomHorizontalFlip(p=0.8),
  transforms.AutoAugment(policy=AutoAugmentPolicy.SVHN),
  transforms.RandomAffine(degrees=60, translate=(0.12, 0.12), scale=(0.8, 1.2)),
	transforms.ToTensor(),
  transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_tfm5 = transforms.Compose([
	transforms.Resize((256, 256)),
  # transforms.Resize((128, 128)),
  # transforms.RandomChoice(randomAugment),
  transforms.ColorJitter(0.3, 0.2, 0.15),
  # transforms.GaussianBlur(5),
  transforms.RandomHorizontalFlip(p=0.6),
  transforms.AutoAugment(policy=AutoAugmentPolicy.IMAGENET),
	transforms.ToTensor(),
  transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

test_tfm = transforms.Compose([
    transforms.Resize((256, 256)),
    # transforms.Resize((128, 128)),
    # transforms.CenterCrop(128),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

"""## Load Data
把所有資料包裝好之後載下來
"""

# Construct datasets.
# The argument "loader" tells how torchvision reads the data.
train_set1 = IMAGE("train/train/", transform=train_tfm1)
train_set2 = IMAGE("train/train/", transform=train_tfm2)
train_set3 = IMAGE("train/train/", transform=train_tfm3)
train_set4 = IMAGE("train/train/", transform=train_tfm4)
train_set5 = IMAGE("train/train/", transform=train_tfm5)
train_set = ConcatDataset([train_set1, train_set2, train_set3, train_set4, train_set5])

unlabeled_set1 = IMAGE("unlabeled_data/", transform=train_tfm1)
unlabeled_set2 = IMAGE("unlabeled_data/", transform=train_tfm2)
unlabeled_set3 = IMAGE("unlabeled_data/", transform=train_tfm3)
unlabeled_set4 = IMAGE("unlabeled_data/", transform=train_tfm4)
unlabeled_set5 = IMAGE("unlabeled_data/", transform=train_tfm5)
unlabeled_set = ConcatDataset([unlabeled_set1, unlabeled_set2, unlabeled_set3, unlabeled_set4, unlabeled_set5])

train_set, validation_set = data.random_split(train_set, [8279*5, 5000])
test_set = IMAGE("test/test/", transform=test_tfm)

print('# images in trainset:', len(train_set)) 
print('# images in validationset:', len(validation_set))
# print('# images in unlabeledset:', len(unlabeled_set))
print('# images in testset:', len(test_set))

# Construct data loaders.
batch_size = 32
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True)
valid_loader = DataLoader(validation_set, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True)
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False)

# get some random training images
dataiter = iter(train_loader)
images, labels = dataiter.next()
print('Image tensor in each batch:', images.shape, images.dtype)
print('Label tensor in each batch:', labels.shape, labels.dtype)

"""看一下要train的圖片，確認做完data augmentation後的圖片合不合理"""

import matplotlib.pyplot as plt
import numpy as np

# functions to show an image
def imshow(img):
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    
# show images
imshow(torchvision.utils.make_grid(images))
# print labels
print('Labels:')
print(' '.join('%5s' % labels[j] for j in range(16)))

"""# **Model**"""

# 看一下哪一種GPU，不過現在免費版的只剩下超爛的K80
use_cuda = torch.cuda.is_available()
torch.manual_seed(87)
device = torch.device("cuda" if use_cuda else "cpu")
print('Device used:', device)
!nvidia-smi

import torchvision.models as models 
model = models.efficientnet_b4(pretrained=True).to(device)
model.classifier = nn.Sequential(
    nn.Dropout(p=0.4, inplace=True),
    nn.Linear(in_features=1792, out_features=4, bias=True),
    ).to(device)

print(model)

"""# **Training**"""

from tqdm.auto import tqdm

class PSEUDO(Dataset):
  def __init__(self, x, y):
      super().__init__()
      self.x = x
      self.y = y
      self.len = len(self.x)
  def __getitem__(self, index):
      return self.x[index], self.y[index]
  def __len__(self):
      return self.len

def get_pseudo_labels(dataset, model, threshold):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
    softmax = nn.Softmax(dim=-1)
    predictions = [[] for i in range(4)]
    Type_rec = np.zeros([4])
    model.eval()

    for batch in tqdm(data_loader):
        img, _ = batch
        with torch.no_grad():
            logits = model(img.to(device))
        probs = softmax(logits)

        for i in range(img.shape[0]):
            _ , Type = torch.max(probs[i], 0)
            Type = int(Type.item())
            if probs[i][Type].item() >= threshold:
                predictions[Type].append(img[i])
                Type_rec[Type] += 1

    ret_predictions = []
    type_list = []

    for i in range(4):
        for j in predictions[i]:
            ret_predictions.append(j)
            type_list.append(i)

    print('Added %d images' %(len(type_list)))
    return PSEUDO(ret_predictions, type_list)

import torchvision.models as models

device = "cuda" if torch.cuda.is_available() else "cpu"
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.0001, weight_decay=1e-5)
n_epochs = 3
best_acc = 0
do_semi = False
threshold = 0.98
# checkpoint = torch.load('drive/MyDrive/model.pt')
# model.load_state_dict(checkpoint['model_state_dict'])
# optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

for epoch in range(n_epochs):
    if do_semi:
        pseudo_set = get_pseudo_labels(unlabeled_set, model, threshold)
        concat_dataset = ConcatDataset([train_set, pseudo_set])
        train_loader = DataLoader(concat_dataset, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=True)
    #start training
    model.train()
    train_loss = []
    train_accs = []
    for batch in tqdm(train_loader):
        imgs, labels = batch
        logits = model(imgs.to(device))
        loss = criterion(logits, labels.to(device))
        # loss = checkpoint['loss']
        optimizer.zero_grad()
        loss.backward()
        grad_norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=10)
        optimizer.step()
        acc = (logits.argmax(dim=-1) == labels.to(device)).float().mean()
        train_loss.append(loss.item())
        train_accs.append(acc)
    train_loss = sum(train_loss) / len(train_loss)
    train_acc = sum(train_accs) / len(train_accs)
    print(f"[ Train | {epoch + 1:03d}/{n_epochs:03d} ] loss = {train_loss:.5f}, acc = {train_acc:.5f}")

    model.eval()
    valid_loss = []
    valid_accs = []
    for batch in tqdm(valid_loader):
        imgs, labels = batch
        with torch.no_grad():
          logits = model(imgs.to(device))
        loss = criterion(logits, labels.to(device))
        acc = (logits.argmax(dim=-1) == labels.to(device)).float().mean()
        valid_loss.append(loss.item())
        valid_accs.append(acc)
    valid_loss = sum(valid_loss) / len(valid_loss)
    valid_acc = sum(valid_accs) / len(valid_accs)
    print(f"[ Valid | {epoch + 1:03d}/{n_epochs:03d} ] loss = {valid_loss:.5f}, acc = {valid_acc:.5f}")
    if valid_acc > best_acc:
        best_acc = valid_acc
        torch.save({
          'optimizer_state_dict': optimizer.state_dict(),
          'model_state_dict': model.state_dict(),
          'loss': loss
          },'model.pt')
        print('model saved')
    # if valid_acc > 0.82:
    #     do_semi = True
    # if valid_acc < 0.7:
    #     do_semi = False

"""# **Method2: 用Torchensemble train**"""

!pip install torchensemble

from torchensemble import VotingClassifier

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f'DEVICE: {device}')
model_path = './model.ckpt'

model = VotingClassifier(estimator=model,n_estimators=3,cuda=True).to(device)
criterion = nn.CrossEntropyLoss() 
model.set_optimizer("Adam",lr=0.0001,weight_decay=1e-05)

# start training
model.fit(train_loader,epochs=2)
torch.save(model.state_dict(), model_path)
print('saving model')

"""# **Prediction**"""

model.eval()
filename_list = []
path = 'test/test/'

for infile in glob.glob(os.path.join(path,'*')):
  filename_list.append(infile[10:])
print(filename_list)
predictions = []

for batch in tqdm(test_loader):
    imgs, labels = batch
    with torch.no_grad():
        logits = model(imgs.to(device))
    predictions.extend(logits.argmax(dim=-1).cpu().numpy().tolist())

with open("predict_ML_final.csv", "w") as f:
    f.write("Name,Type\n")
    for i, pred in  enumerate(predictions):
         f.write(f"{filename_list[i]},{pred}\n")

from google.colab import files
files.download("predict_ML_final.csv")