{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Train a model to predict rotation of an image: 0°, 90°, 180° or 270°\n\nIn this simple example I will show to traina  simple model to check the angle of a rotated image. The accuracy can be very high and the model effective to clean the dataset","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport sys\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom datetime import datetime\nfrom tqdm.notebook import tqdm\n\nimport multiprocessing as mp\n\nimport torch\nimport torch.nn as nn\nimport torchvision.transforms as T\nimport torchvision.datasets as datasets\nfrom torch.utils.data import DataLoader\n\nsys.path.insert(1, \"../input/timm-pytorch-image-models/pytorch-image-models-master/\")\n\nimport timm\n\nROOT = '/kaggle/input/hotel-id-2021-fgvc8'\nTRAIN_PATH = os.path.join(ROOT, \"train_images\")\nTEST_PATH = os.path.join(ROOT, \"test_images\")\nTRAIN_CSV = os.path.join(ROOT, \"train.csv\")\n\n# If you want to resize tehe dataset to speed up all the training set it to True\nRESIZE_ALL_DATASET = False\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# load train dataframe\ndf_hotels = pd.read_csv(TRAIN_CSV)\n\n# Don't sample it ! It's just to save time\ndf_hotels = df_hotels.sample(1000)\nprint(f\"len df_hotels = {len(df_hotels)}\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 1. Resize the dataset (maximum width or height: 512)\n\nResize the dataset will help us to speed up the dataloader and the also the entire training","metadata":{}},{"cell_type":"code","source":"if RESIZE_ALL_DATASET:\n    import threading\n\n    # Use multithreading to speed up the resize job\n    def resize_thread(df, t_id):\n        n_tot = len(df)\n        for item in tqdm(df.iterrows(), total=n_tot):\n            file_path = os.path.join(TRAIN_PATH, str(item[1][\"chain\"]), item[1][\"image\"])\n            img = Image.open(file_path)\n            width, height = img.size\n            if width > 512 or height > 512:\n                max_size = max(width, height)\n                scale = 512/max_size\n                img = img.resize((int(width*scale), int(scale*height)))\n                out_folder = os.path.join(\"train_resized\", str(item[1][\"chain\"]))\n                os.makedirs(out_folder, exist_ok=True)\n                img.save(os.path.join(out_folder, item[1][\"image\"]))\n\n    image_name = df_hotels[\"image\"].values\n    chain = df_hotels[\"chain\"].values\n\n    # create the new resized dataset dir\n    os.makedirs(\"train_resized\", exist_ok=True)\n\n    n_tot = len(df_hotels)\n    df_hotels.index = range(n_tot)\n\n    # Start threads\n    num_thread = mp.cpu_count()\n    num_elems = n_tot//num_thread\n    thread_list = []\n    for i in range(num_thread):\n        i1 = i*num_elems\n        if i==num_thread-1:\n            i2 = n_tot\n        else:\n            i2 = (i+1)*num_elems\n\n        thread_list += [threading.Thread(target=resize_thread, args=(df_hotels.iloc[i1:i2],i))]\n\n    for i in range(num_thread):\n        thread_list[i].start()\n\n    for i in range(num_thread):\n        thread_list[i].join()\n    \n    TRAIN_PATH = \"train_resized\"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 1.1 Compute the aspect ratio of images","metadata":{}},{"cell_type":"code","source":"image_name = df_hotels[\"image\"].values\nchain = df_hotels[\"chain\"].values\n# Add the information of image width, height and aspect ration to the dataframe\nwidth = []\nheight = []\nfor i in tqdm(range(len(image_name))):\n    img = Image.open(os.path.join(TRAIN_PATH, str(chain[i]), image_name[i]))\n    w, h = img.size\n    width += [w]\n    height += [h]\nwidth = np.array(width)\nheight = np.array(height)\n\ndf_hotels[\"width\"] = width\ndf_hotels[\"height\"] = height\ndf_hotels[\"ar\"] = width/height\ndf_hotels.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 1.2 Check images with different aspect ratio to check if rotations happen more frequently on a specific AR","metadata":{}},{"cell_type":"code","source":"def show_ar_images(index):\n    fig, ax = plt.subplots(5,10, figsize=(30,15))\n    for i in range(50):\n        sample = df_hotels[index].sample(1)\n        img = Image.open(os.path.join(TRAIN_PATH, str(sample[\"chain\"].item()), sample[\"image\"].item()))\n        ax[i//10,i%10].imshow(img)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_ar_images(df_hotels[\"ar\"]<=0.75)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_ar_images(df_hotels[\"ar\"]>0.75)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"It seems that rotations are more frequent for aspect ratio greather than 0.75, then we can train our network using images with ar <= 0.75","metadata":{}},{"cell_type":"markdown","source":"# 2. Training","metadata":{}},{"cell_type":"code","source":"import sklearn.metrics as metrics\n# Hyperparameters\n# there are 4 possible rotations: [0, 90, 180, 270]  degrees\nnum_classes = 4\nbatch_size = 64\n# number of rotations to sample from [0, 90, 180, 270] during training\nnum_rotations = 2\n# number of workers for the data loader\nnum_workers = mp.cpu_count()\n# number of epochs ----------------------- (CHANGE to 40 !)\nepochs = 10\n# learning rate\nlr = 5e-4\n#  label smoothing eps\nsmooth_eps = 0.15","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n## 2.1 Splitting\nFirst of all split the dataset to train and validation set","metadata":{}},{"cell_type":"code","source":"# Dataset splitting: 85 % train and 15 % validation\ndf_hotels_ar = df_hotels[df_hotels[\"ar\"]<=0.75]\nn_sample = len(df_hotels_ar)\nn_train = int(n_sample*0.85)\nindexes = np.arange(n_sample)\nnp.random.shuffle(indexes)\ndf_train = df_hotels_ar.iloc[indexes[:n_train],:]\ndf_val = df_hotels_ar.iloc[indexes[n_train:],:]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.2 Define the model","metadata":{}},{"cell_type":"code","source":"# Load pretrained model (on imagenet)\nmodel = timm.create_model(\"efficientnet_b0\", pretrained=False, checkpoint_path=\"../input/timm-pretrained-efficientnet/efficientnet/efficientnet_b0_ra-3dd342df.pth\",drop_rate=0.4)\n# Change the number of classes\nmodel.classifier = nn.Linear(model.bn2.num_features, num_classes)\n# get the default config\nconfig = timm.data.resolve_data_config({}, model=model)\n# get the default input size for the defined model\ninput_size = config[\"input_size\"][1:]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.3 Create datasets, preprocessing and dataloader","metadata":{}},{"cell_type":"code","source":"class RotationDataset(datasets.vision.VisionDataset):\n    def __init__(\n        self,\n        root,\n        df_hotels,        # The dataframe\n        num_rotations=2,  # how many rotations of the same image choosen randomly by [0, 90, 180, 270]\n        validation=False, # When validate the dataset, we want to validate all the rotations\n        transform=None,\n        loader=datasets.folder.default_loader,\n    ):\n        super(RotationDataset, self).__init__(\n            root, transform=transform, target_transform=None\n        )\n        self.loader = loader\n        self.validation = validation\n        self.num_rotations = num_rotations\n\n        image_name = df_hotels[\"image\"].values\n        chain = df_hotels[\"chain\"].values\n\n        self.samples = []\n        for i in range(len(image_name)):\n            self.samples.append(os.path.join(root, str(chain[i]), image_name[i]))\n\n        self.samples = np.array(self.samples)\n\n    def __getitem__(self, index):\n        sample = self.loader(self.samples[index])\n\n        if self.transform is not None:\n            sample = self.transform(sample)\n\n        samples = []\n\n        if self.validation:\n            labels = [0, 1, 2, 3]\n        else:\n            np.random.seed(index)\n            labels = np.random.choice([0, 1, 2, 3], self.num_rotations)\n        for l in labels:\n            if l == 0:\n                samples += [sample]\n            elif l == 1:\n                samples += [sample.transpose(1, 2)]   # This is equivalent of rotate(90)\n            elif l == 2:\n                samples += [sample.flip(1)]           # This is equivalent of rotate(180)\n            elif l == 3:\n                samples += [sample.transpose(1, 2).flip(2)]   # This is equivalent of rotate(270)\n\n        return samples, list(labels)\n\n    def __len__(self):\n        return len(self.samples)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Define dataloader and dataset for training","metadata":{}},{"cell_type":"code","source":"# Very simple data augmentation for the training dataset\ntrain_transform = T.Compose([T.RandomRotation(15),\n                            T.RandomResizedCrop(input_size, scale=(0.5,1.0)),\n                            T.RandomHorizontalFlip(0.5),\n                            T.ColorJitter(0.15, 0.15, 0.15, 0.1),\n                            T.RandomGrayscale(0.2),\n                            T.ToTensor(),\n                            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])\n# No augmentation for validation dataset\nval_transform = T.Compose([T.Resize(input_size),\n                            T.ToTensor(),\n                            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])\n\ntrain_dataset = RotationDataset(TRAIN_PATH, df_train, num_rotations=num_rotations, transform=train_transform)\nval_dataset = RotationDataset(TRAIN_PATH, df_val, validation=True, transform=val_transform)\n\n# Data loaders\ntrain_loader = DataLoader(\n    train_dataset,\n    batch_size=batch_size,\n    shuffle=True,\n    num_workers=num_workers,\n    pin_memory=True,\n)\n\nval_loader = DataLoader(\n    val_dataset,\n    batch_size=batch_size,\n    shuffle=False,\n    num_workers=num_workers,\n    pin_memory=True,\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.4 Define optimizer and schedulers","metadata":{}},{"cell_type":"code","source":"# Label smoothing cross entropy loss\nclass LabelSmoothingCrossEntropy(nn.Module):\n    def __init__(self, classes, epsilon=0.2, dim=-1):\n        super(LabelSmoothingCrossEntropy, self).__init__()\n        self.confidence = 1.0 - epsilon\n        self.epsilon = epsilon\n        self.cls = classes\n        self.dim = dim\n\n    def forward(self, pred, target):\n        pred = pred.log_softmax(dim=self.dim)\n        with torch.no_grad():\n            # true_dist = pred.data.clone()\n            true_dist = torch.zeros_like(pred)\n            true_dist.fill_(self.epsilon / (self.cls - 1))\n            true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)\n        return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n# you can also choose SGD\n#optimizer = torch.optim.SGD(model.parameters(), lr=lr) \n\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n    optimizer, T_max=epochs, eta_min=1e-5\n)\n# Cross entropy with label smoothing criterion (it hepls to calibrate the probabilities)\ncriterion = LabelSmoothingCrossEntropy(\n    classes=num_classes,\n    epsilon=smooth_eps,\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.5 Train!","metadata":{}},{"cell_type":"code","source":"# Eval function\ndef eval(\n    model,  data_loader, device=\"cuda:0\"\n):\n    metric_dict = {}\n    pred_list = []\n    label_list = []\n    model.eval()\n    model.to(device)\n    with torch.no_grad():\n        for images, labels in data_loader:\n            \n            images = torch.cat(images).to(device)\n            labels = torch.cat(labels).to(device)\n\n            logits = model(images.cuda())\n            pred_labels = np.argmax(logits.cpu().numpy(), axis=1)\n            pred_list += list(pred_labels)\n            label_list += list(labels.cpu().numpy())\n\n        metric_dict[\"accuracy\"] = metrics.accuracy_score(label_list, pred_list)\n        metric_dict[\"macro_precision\"] = metrics.precision_score(\n            label_list, pred_list, average=\"macro\"\n        )\n        metric_dict[\"macro_recall\"] = metrics.recall_score(\n            label_list, pred_list, average=\"macro\"\n        )\n        metric_dict[\"macro_f1\"] = metrics.f1_score(\n            label_list, pred_list, average=\"macro\"\n        )\n\n        print(\n            f'Val acc. {metric_dict[\"accuracy\"]*100} %, meanPrecision {metric_dict[\"macro_precision\"]}, meanRecall {metric_dict[\"macro_recall\"]}, meanF1 {metric_dict[\"macro_f1\"]}'\n        )\n\n    return metric_dict","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Train function\ndef train(\n    model,\n    optimizer,\n    scheduler,\n    criterion,\n    train_data_loader,\n    val_data_loader,\n    epochs=50,\n    batch_size=32,\n    log_freq=0.25,\n    save_every=10,\n    device=\"cuda:0\",\n):\n    iteration = 1\n    date_now = datetime.now().strftime(\"%Y-%m-%d_%H-%M-%S\")\n\n    num_samples = len(train_data_loader.dataset)\n    num_batches = num_samples // batch_size\n    log_step = max(int(log_freq * num_batches), 1)\n\n    model.to(device)\n    for epoch in range(1, epochs + 1):\n        model.train()\n        for batch, (images, labels) in enumerate(train_data_loader):\n            optimizer.zero_grad()\n\n            images = torch.cat(images).to(device)\n            labels = torch.cat(labels).to(device)\n\n            logits = model(images)\n\n            loss = criterion(logits, labels)\n            loss.backward()\n\n            optimizer.step()\n\n            iteration += 1\n            if batch % log_step == 0:\n                print(\n                    f\"[{batch}/{num_batches}] Epoch {epoch} : Train loss {loss.item()}\"\n                )\n\n        scheduler.step()\n\n        metric_dict = eval(model, val_data_loader, device=device)\n\n        # Save checkpoints\n        if epoch % save_every == 0:\n            pth_name = (\n                f'{epoch:03d}_{100*metric_dict[\"accuracy\"]:.4f}.pth'\n            )\n            print(f\"Saving {pth_name}\")\n\n            if not os.path.isdir(os.path.join(\"checkpoints\", date_now)):\n                os.makedirs(os.path.join(\"checkpoints\", date_now), exist_ok=True)\n\n            torch.save(\n                model.state_dict(), os.path.join(\"checkpoints\", date_now, pth_name)\n            )","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"    train(\n        model,\n        optimizer,\n        scheduler,\n        criterion,\n        train_loader,\n        val_loader,\n        epochs=epochs,\n        batch_size=batch_size,\n        device=device,\n    )","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}