<filename>core/data/dataset.py
from __future__ import print_function

import torch
import torch.utils.data
import torchvision.transforms as transforms

import numpy as np
import argparse
import os
from PIL import Image


class MetaMSTAR(torch.utils.data.Dataset):
    """MSTAR Dataset for meta-based few-shot SAR image classification.

    Args:
        root (string): Root directory of dataset where ``MSTAR`` exists.
        train (bool, optional): If True, creates dataset from ``meta_train``,
            otherwise from ``meta_test``.
        transform (callable, optional): A function/transform that takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
    """

    training_folder = 'meta_train'
    training_categories = ['BMP2', 'BTR60', 'BTR70', 'D7', 'T62', 'T72', 'ZIL131']
    test_folder = 'meta_test'
    test_categories = ['2S1', 'BRDM2', 'ZSU234']

    def __init__(
            self,
            args,
            root,
            train=True,
            transform=None,
            method_type=None
    ):
        super(MetaMSTAR, self).__init__()
        self.method_type = method_type
        self.root = root
        self.train = train
        if train:
            self.num_tasks = args['train_episode']
            self.num_ways = args['train_way']
            self.num_shots = args['train_shot']
            self.num_queries = args['train_query']
        else:
            self.num_tasks = args['test_episode']
            self.num_ways = args['test_way']
            self.num_shots = args['test_shot']
            self.num_queries = args['test_query']
        self.depression_angle = args['depression_angle']

        # transform need to be defined outside
        assert transform is not None
        self.transform = transform

        if self.train:
            if method_type == 'preTrain':
                pass
            else:
                self.num_categories = len(self.training_categories)
                self.data_dir = os.path.join(self.root, self.training_folder)
                self.lab2img_paths = dict()
                for i, category in enumerate(self.training_categories):
                    img_dir = os.path.join(self.data_dir, category, str(self.depression_angle))
                    for _, _, files in os.walk(img_dir):
                        paths = [os.path.join(img_dir, file) for file in files]
                        self.lab2img_paths[i] = paths
        else:
            self.num_categories = len(self.test_categories)
            self.data_dir = os.path.join(self.root, self.test_folder)
            self.lab2img_paths = dict()
            for i, category in enumerate(self.test_categories):
                img_dir = os.path.join(self.data_dir, category, str(self.depression_angle))
                for _, _, files in os.walk(img_dir):
                    paths = [os.path.join(img_dir, file) for file in files]
                    self.lab2img_paths[i] = paths

    def __getitem__(self, index):
        if self.method_type == 'preTrain':
            pass
        else:
            # np.random.seed(index)
            support_xs = []
            support_ys = []
            query_xs = []
            query_ys = []
            # sample class index for the task
            cat_sampled = np.random.choice(self.num_categories, self.num_ways, replace=False)
            for i, cat in enumerate(cat_sampled):
                # sample image index for the support set and the query set
                img_paths = self.lab2img_paths[cat]
                indices_sampled = np.random.choice(
                    len(img_paths), self.num_shots + self.num_queries, replace=False)
                support_ids, query_ids = np.split(indices_sampled, [self.num_shots])
                # construct the support set
                for j in support_ids:
                    img = Image.open(img_paths[j])
                    img = self.transform(img)
                    support_xs.append(img)
                    support_ys.append(i)
                # construct the query set
                for j in query_ids:
                    img = Image.open(img_paths[j])
                    img = self.transform(img)
                    query_xs.append(img)
                    query_ys.append(i)
            support_xs, support_ys = torch.stack(support_xs).repeat(1, 3, 1, 1), torch.tensor(support_ys)
            query_xs, query_ys = torch.stack(query_xs).repeat(1, 3, 1, 1), torch.tensor(query_ys)
            image_size = support_xs.size(2)
            support_xs = support_xs.view(self.num_ways, self.num_shots, 3, image_size, image_size)
            query_xs = query_xs.view(self.num_ways, self.num_queries, 3, image_size, image_size)
            x = torch.cat([support_xs, query_xs], dim=1)
            return x

    def __len__(self):
        if self.method_type == 'preTrain':
            pass
        else:
            return self.num_tasks

<filename>core/data/dataset.py
from __future__ import print_function

import torch
import torch.utils.data
import torchvision.transforms as transforms

import numpy as np
import argparse
import os
from PIL import Image


class MetaMSTAR(torch.utils.data.Dataset):
    """MSTAR Dataset for meta-based few-shot SAR image classification.

    Args:
        root (string): Root directory of dataset where ``MSTAR`` exists.
        train (bool, optional): If True, creates dataset from ``meta_train``,
            otherwise from ``meta_test``.
        transform (callable, optional): A function/transform that takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
    """

    training_folder = 'meta_train'
    training_categories = ['BMP2', 'BTR60', 'BTR70', 'D7', 'T62', 'T72', 'ZIL131']
    test_folder = 'meta_test'
    test_categories = ['2S1', 'BRDM2', 'ZSU234']

    def __init__(
            self,
            args,
            root,
            train=True,
            transform=None,
            method_type=None
    ):
        super(MetaMSTAR, self).__init__()
        self.method_type = method_type
        self.root = root
        self.train = train
        if train:
            self.num_tasks = args['train_episode']
            self.num_ways = args['train_way']
            self.num_shots = args['train_shot']
            self.num_queries = args['train_query']
        else:
            self.num_tasks = args['test_episode']
            self.num_ways = args['test_way']
            self.num_shots = args['test_shot']
            self.num_queries = args['test_query']
        self.depression_angle = args['depression_angle']

        # transform need to be defined outside
        assert transform is not None
        self.transform = transform

        if self.train:
            if method_type == 'preTrain':
                pass
            else:
                self.num_categories = len(self.training_categories)
                self.data_dir = os.path.join(self.root, self.training_folder)
                self.lab2img_paths = dict()
                for i, category in enumerate(self.training_categories):
                    img_dir = os.path.join(self.data_dir, category, str(self.depression_angle))
                    for _, _, files in os.walk(img_dir):
                        paths = [os.path.join(img_dir, file) for file in files]
                        self.lab2img_paths[i] = paths
        else:
            self.num_categories = len(self.test_categories)
            self.data_dir = os.path.join(self.root, self.test_folder)
            self.lab2img_paths = dict()
            for i, category in enumerate(self.test_categories):
                img_dir = os.path.join(self.data_dir, category, str(self.depression_angle))
                for _, _, files in os.walk(img_dir):
                    paths = [os.path.join(img_dir, file) for file in files]
                    self.lab2img_paths[i] = paths

    def __getitem__(self, index):
        if self.method_type == 'preTrain':
            pass
        else:
            # np.random.seed(index)
            support_xs = []
            support_ys = []
            query_xs = []
            query_ys = []
            # sample class index for the task
            cat_sampled = np.random.choice(self.num_categories, self.num_ways, replace=False)
            for i, cat in enumerate(cat_sampled):
                # sample image index for the support set and the query set
                img_paths = self.lab2img_paths[cat]
                indices_sampled = np.random.choice(
                    len(img_paths), self.num_shots + self.num_queries, replace=False)
                support_ids, query_ids = np.split(indices_sampled, [self.num_shots])
                # construct the support set
                for j in support_ids:
                    img = Image.open(img_paths[j])
                    img = self.transform(img)
                    support_xs.append(img)
                    support_ys.append(i)
                # construct the query set
                for j in query_ids:
                    img = Image.open(img_paths[j])
                    img = self.transform(img)
                    query_xs.append(img)
                    query_ys.append(i)
            support_xs, support_ys = torch.stack(support_xs).repeat(1, 3, 1, 1), torch.tensor(support_ys)
            query_xs, query_ys = torch.stack(query_xs).repeat(1, 3, 1, 1), torch.tensor(query_ys)
            image_size = support_xs.size(2)
            support_xs = support_xs.view(self.num_ways, self.num_shots, 3, image_size, image_size)
            query_xs = query_xs.view(self.num_ways, self.num_queries, 3, image_size, image_size)
            x = torch.cat([support_xs, query_xs], dim=1)
            return x

    def __len__(self):
        if self.method_type == 'preTrain':
            pass
        else:
            return self.num_tasks

MSTAR Dataset for meta-based few-shot SAR image classification.      Args:         root (string): Root directory of dataset where ``MSTAR`` exists.         train (bool, optional): If True, creates dataset from ``meta_train``,             otherwise from ``meta_test``.         transform (callable, optional): A function/transform that takes in an PIL image             and returns a transformed version. E.g, ``transforms.RandomCrop`` # transform need to be defined outside # np.random.seed(index) # sample class index for the task # sample image index for the support set and the query set # construct the support set # construct the query set