core/dsproc_mclsmfolder.py

import os
import torch
from PIL import Image
from collections import OrderedDict
from toolkit.dhelper import traverse_recursively
import random
from torch import nn
import numpy as np
import timm
import einops
import torch.nn.functional as F


class SRMConv2d_simple(nn.Module):
    def __init__(self, inc=3):
        super(SRMConv2d_simple, self).__init__()
        self.truc = nn.Hardtanh(-3, 3)
        self.kernel = torch.from_numpy(self._build_kernel(inc)).float()

    def forward(self, x):
        out = F.conv2d(x, self.kernel, stride=1, padding=2)
        out = self.truc(out)

        return out

    def _build_kernel(self, inc):
        # filter1: KB
        filter1 = [[0, 0, 0, 0, 0],
                   [0, -1, 2, -1, 0],
                   [0, 2, -4, 2, 0],
                   [0, -1, 2, -1, 0],
                   [0, 0, 0, 0, 0]]
        # filter2：KV
        filter2 = [[-1, 2, -2, 2, -1],
                   [2, -6, 8, -6, 2],
                   [-2, 8, -12, 8, -2],
                   [2, -6, 8, -6, 2],
                   [-1, 2, -2, 2, -1]]
        # filter3：hor 2rd
        filter3 = [[0, 0, 0, 0, 0],
                   [0, 0, 0, 0, 0],
                   [0, 1, -2, 1, 0],
                   [0, 0, 0, 0, 0],
                   [0, 0, 0, 0, 0]]

        filter1 = np.asarray(filter1, dtype=float) / 4.
        filter2 = np.asarray(filter2, dtype=float) / 12.
        filter3 = np.asarray(filter3, dtype=float) / 2.
        # statck the filters
        filters = [[filter1],  # , filter1, filter1],
                   [filter2],  # , filter2, filter2],
                   [filter3]]  # , filter3, filter3]]
        filters = np.array(filters)
        filters = np.repeat(filters, inc, axis=1)
        return filters


class MultiClassificationProcessor_mfolder(torch.utils.data.Dataset):
    def __init__(self, transform=None):
        self.transformer_  = transform
        self.extension_    = '.jpg .jpeg .png .bmp .webp .tif .eps'
        # load category info
        self.ctg_names_    = []             # ctg_idx to ctg_name     
        self.ctg_name2idx_ = OrderedDict()  # ctg_name to ctg_idx
        # load image infos
        self.img_names_  = []    # img_idx to img_name
        self.img_paths_  = []    # img_idx to img_path
        self.img_labels_ = []    # img_idx to img_label

        self.srm = SRMConv2d_simple()

    def load_data_from_dir_test(self, folders):

        # Load image from folder.

        # Args:
        #     dataset_list: dictionary where key is a label and value is a list of folder paths.
        print(folders)
        img_paths = []
        traverse_recursively(folders, img_paths, self.extension_)

        for img_path in img_paths:
            img_name = os.path.basename(img_path)
            self.img_names_.append(img_name)
            self.img_paths_.append(img_path)

        length = len(img_paths)
        print('log: {} image num is {}'.format(folders, length))

    def load_data_from_dir(self, dataset_list):

        # Load image from folder.

        # Args:
        #     dataset_list: dictionary where key is a label and value is a list of folder paths.

        for ctg_name, folders in dataset_list.items():

            if ctg_name not in self.ctg_name2idx_:
                self.ctg_name2idx_[ctg_name] = len(self.ctg_names_)
                self.ctg_names_.append(ctg_name)

            for img_root in folders:
                img_paths = []
                traverse_recursively(img_root, img_paths, self.extension_)

                print(img_root)

                length = len(img_paths)
                for i in range(length):
                    img_path = img_paths[i]
                    img_name = os.path.basename(img_path)
                    self.img_names_.append(img_name)
                    self.img_paths_.append(img_path)
                    self.img_labels_.append(self.ctg_name2idx_[ctg_name])

            print('log: category is %d(%s), image num is %d' % (self.ctg_name2idx_[ctg_name], ctg_name, length))

    def load_data_from_txt(self, img_list_txt, ctg_list_txt):
        """Load image from txt.

        Args:
            img_list_txt: image txt, format is [file_path, ctg_idx].
            ctg_list_txt: category txt, format is [ctg_name, ctg_idx].
        """
        # check
        assert os.path.exists(img_list_txt), 'log: does not exist: {}'.format(img_list_txt)
        assert os.path.exists(ctg_list_txt), 'log: does not exist: {}'.format(ctg_list_txt)

        # load category
        # : open category info file
        with open(ctg_list_txt) as f:
            ctg_infos = [line.strip() for line in f.readlines()] 
        # :load category name & category index
        for ctg_info in ctg_infos:
            tmp      = ctg_info.split(' ')
            ctg_name = tmp[0]
            ctg_idx  = int(tmp[1])
            self.ctg_name2idx_[ctg_name] = ctg_idx
            self.ctg_names_.append(ctg_name)

        # load sample
        # : open image info file
        with open(img_list_txt) as f:
            img_infos = [line.strip() for line in f.readlines()]
        random.shuffle(img_infos)
        # : load image path & category index
        for img_info in img_infos:
            img_path, ctg_name = img_info.rsplit(' ', 1)
            img_name = img_path.split('/')[-1]
            ctg_idx  = int(ctg_name)
            self.img_names_.append(img_name)
            self.img_paths_.append(img_path)
            self.img_labels_.append(ctg_idx)

        for ctg_name in self.ctg_names_:
            print('log: category is %d(%s), image num is %d' % (self.ctg_name2idx_[ctg_name], ctg_name, self.img_labels_.count(self.ctg_name2idx_[ctg_name])))

    def _add_new_channels_worker(self, image):
        new_channels = []

        image = einops.rearrange(image, "h w c -> c h w")
        image = (image- torch.as_tensor(timm.data.constants.IMAGENET_DEFAULT_MEAN).view(-1, 1, 1)) / torch.as_tensor(timm.data.constants.IMAGENET_DEFAULT_STD).view(-1, 1, 1)
        srm = self.srm(image.unsqueeze(0)).squeeze(0)
        new_channels.append(einops.rearrange(srm, "c h w -> h w c").numpy())

        new_channels = np.concatenate(new_channels, axis=2)
        return torch.from_numpy(new_channels).float()

    def add_new_channels(self, images):
        images_copied = einops.rearrange(images, "c h w -> h w c")
        new_channels = self._add_new_channels_worker(images_copied)
        images_copied = torch.concatenate([images_copied, new_channels], dim=-1)
        images_copied = einops.rearrange(images_copied, "h w c -> c h w")

        return images_copied

    def __getitem__(self, index):
        img_path = self.img_paths_[index]

        img_data = Image.open(img_path).convert('RGB')
        img_size = img_data.size[::-1]   # [h, w]

        all_data = []
        for transform in self.transformer_:
            current_data = transform(img_data)
            current_data = self.add_new_channels(current_data)
            all_data.append(current_data)
        img_label = self.img_labels_[index]

        return torch.stack(all_data, dim=0), img_label, img_path, img_size

    def __len__(self):
        return len(self.img_names_)