# Uncomment and run the appropriate command for your operating system, if required

# Linux / Binder / Windows (No GPU)
# !pip install numpy matplotlib torch==1.7.0+cpu torchvision==0.8.1+cpu torchaudio==0.7.0 -f https://download.pytorch.org/whl/torch_stable.html

# Linux / Windows (GPU)
# pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
 
# MacOS (NO GPU)
# !pip install numpy matplotlib torch torchvision torchaudio

!pip install opendatasets --upgrade

import os
import opendatasets as od
import numpy as np

import torch
import torchvision
from torch.utils.data import random_split
from torch.utils.data.dataloader import DataLoader
import torch.nn as nn
import torch.nn.functional as F
from torchvision.datasets.utils import download_url
from torchvision.datasets import ImageFolder
from torchvision.transforms import ToTensor
from torchvision.utils import make_grid

import matplotlib
import matplotlib.pyplot as plt
%matplotlib inline

matplotlib.rcParams['figure.facecolor'] = '#ffffff'

project_name='zerotogans-project-intel-image-classification'

dataset_url = 'https://www.kaggle.com/puneet6060/intel-image-classification'
od.download(dataset_url)

dataset_folder = './intel-image-classification'
tarin_folder = '/seg_train/seg_train'
test_folder = '/seg_test/seg_test'
pred_folder = '/seg_pred/seg_pred'

dataset = ImageFolder(dataset_folder+tarin_folder, transform=ToTensor())
test_ds = ImageFolder(dataset_folder+test_folder, transform=ToTensor())
classes = dataset.classes

def describe_dataset(data):
    classes_count = { c:0 for c in classes}
    for _,label in data:
        classes_count[classes[label]]+=1

    print(f'Classes: {len(classes)} - {classes} ')
    print(f'Examples: {len(data)}')
    print(f'Counts: {classes_count}')
    plt.bar(classes_count.keys(),classes_count.values(), alpha=0.3 )

    return classes_count

_ = describe_dataset(dataset)

dataset = list(filter(lambda x: x[0].shape == torch.Size([3, 150, 150]), dataset))
test_ds = list(filter(lambda x: x[0].shape == torch.Size([3, 150, 150]), test_ds))

classes_count = describe_dataset(dataset)

min_count = min(classes_count.values())

new_dataset = []
excedent_count = {k: v - min_count for k, v in classes_count.items()}
for item in dataset:
    _, label = item
    if excedent_count[classes[label]]>0:
        excedent_count[classes[label]] -=1
    else:
        new_dataset.append(item)
dataset = new_dataset

classes_count = describe_dataset(new_dataset)

random_seed = 42
torch.manual_seed(random_seed);

val_size = 3000
train_size = len(dataset) - val_size

train_ds, val_ds = random_split(dataset, [train_size, val_size])

print("\nTRAIN DATASET")
_=describe_dataset(train_ds)

print("\nVALIDATION DATASET")
_=describe_dataset(val_ds)

print("\nTEST DATASET")
_=describe_dataset(test_ds)

dataset = None
new_dataset = None

jovian.log_dataset(dataset_url=dataset_url, val_size=val_size, random_seed=random_seed)

print(train_ds[0][0].shape)
train_ds[0]

def show_example(img, label):
    print('Label: ', classes[label], "("+str(label)+")")
    plt.imshow(img.permute(1, 2, 0))

show_example(*train_ds[0])

show_example(*train_ds[10000])

batch_size=256

train_dl = DataLoader(train_ds, batch_size, shuffle=True, num_workers=3, pin_memory=True)
val_dl = DataLoader(val_ds, batch_size*2, num_workers=3, pin_memory=True)
test_dl = DataLoader(test_ds, batch_size*2, num_workers=3, pin_memory=True)

def show_batch(dl):
    for images, labels in dl:
        fig, ax = plt.subplots(figsize=(24, 12))
        ax.set_xticks([]); ax.set_yticks([])
        ax.imshow(make_grid(images, nrow=16).permute(1, 2, 0))
        break

show_batch(train_dl)