dataset.py

from pathlib import Path
from tqdm import tqdm
from PIL import Image
import numpy as np
import lz4.frame
import argparse
import pickle
import random
import torch
import os
from collections import defaultdict
from typing import Union, List, Any
from torchvision import transforms

def load_sequences(sequences_file: str):
    with open(sequences_file, 'rb') as f:
        source_folder, class_names, sequences = pickle.load(f)
    return source_folder, class_names, sequences


def sequences_info(class_names, sequences):
    # Information on the loaded sequences
    print(f'Number of sequences: {len(sequences)}')
    classes = defaultdict(list)
    for seqid, seq in sequences.items():
        classes[seq['class']].append(seqid)
    for idx, ids in classes.items():
        print(f'Class {idx}.{class_names[idx]}: {len(ids)}')
    return classes


def cmd_make_sequences(args):
    print('*** make_sequences command ***')
    source_folder = Path(args.source_folder)
    curclass_id = 0
    class_names = []
    sequences = {}
    for dir in tqdm(sorted(source_folder.iterdir())):
        if dir.is_dir():
            # Dictionary of classes and their ids
            class_names.append(dir.name)
            for subfolder in sorted(dir.iterdir()):
                for img in sorted(subfolder.iterdir()):
                    id = img.stem.split("-")[-1]  # Keep last part after "-"
                    id = dir.name + id[:-2]  # Remove sequence index
                    seq = sequences.setdefault(id, {'class': curclass_id, 'images': []})
                    seq['images'].append(img.relative_to(source_folder))
            curclass_id += 1
    print(f'Saving {args.sequences_filename}...', end='')
    with open(args.sequences_filename, 'wb') as f:
        pickle.dump((args.source_folder, class_names, sequences), f)
    print(' done.')
    sequences_info(class_names, sequences)


def cmd_make_split(args):
    print('*** make_split command ***')
    _, class_names, sequences = load_sequences(args.sequences_filename)
    classes = sequences_info(class_names, sequences)
    if args.numtrain:
        print(f'Splitting with {args.numtrain} training samples')
        ids = list(sequences.keys())
        random.shuffle(ids)
        train_ids = ids[:args.numtrain]
        test_ids = ids[args.numtrain:]
    else:
        print(f'Splitting with {args.perctrain}% training samples')
        train_ids = []
        test_ids = []
        for ids in classes.values():
            numtrain = len(ids) * args.perctrain // 100
            random.shuffle(ids)
            train_ids.extend(ids[:numtrain])
            test_ids.extend(ids[numtrain:])

    def classes_stat(class_names, sequences, ids):
        classes = defaultdict(int)
        for id in ids:
            seq = sequences[id]
            classes[seq['class']] += 1
        for idx, count in classes.items():
            print(f'\tClass {idx}.{class_names[idx]}: {count}')

    print(f'Number of training samples: {len(train_ids)}')
    classes_stat(class_names, sequences, train_ids)
    print(f'Number of test samples: {len(test_ids)}')
    classes_stat(class_names, sequences, test_ids)

    # Save ids in file
    print(f'Saving {args.split_filename}...', end='')
    with open(args.split_filename, 'wb') as f:
        pickle.dump((args.sequences_filename, train_ids, test_ids), f)
    print(' done.')


def load_split(split_file: str):
    with open(split_file, 'rb') as f:
        sequences_filename, train_ids, test_ids = pickle.load(f)
    source_folder, class_names, sequences = load_sequences(sequences_filename)
    return source_folder, class_names, sequences, train_ids, test_ids


def tensors_stats(tensors_folder):
     # File containing mean and std
    info_file_path = os.path.join(tensors_folder, 'info.pkl')

    with open(info_file_path, 'rb') as info:
        class_names, mean, stddev = pickle.load(info)
        print('mean: ', mean)
        print('std: ', stddev)

    return class_names, mean, stddev    


def cmd_make_tensors(args):
    print('*** make_tensors command ***')
    source_folder, class_names, sequences, train_ids, test_ids = load_split(args.split_filename)
    source_folder = Path(source_folder)

    to_tensor = transforms.PILToTensor()
    sum_x = np.zeros(3)
    sum_x2 = np.zeros(3)
    num_pixels = 0

    def make_tensors(ids: list[str], subfolder: str):
        nonlocal sum_x, sum_x2, num_pixels
        save_cropped = ''
        cropped_folder = None
        if args.cropped_folder is not None:
            cropped_folder = Path(args.cropped_folder) / subfolder
            cropped_folder.mkdir(parents=True, exist_ok=True)
            save_cropped = ids[0]
        tensors_folder = Path(args.tensors_folder) / subfolder
        tensors_folder.mkdir(parents=True, exist_ok=True)
        for seqid in tqdm(ids):
            s = sequences[seqid]
            class_value = s['class']
            frames = s['images']
            images = []
            for frame in frames:
                frame = source_folder / frame
                frame = Image.open(frame)
                frame = frame.crop(args.crop_box)

                arr = np.array(frame, dtype=float)
                sum_x += np.sum(arr, axis=(0,1))
                sum_x2 += np.sum(arr**2, axis=(0,1))
                num_pixels += frame.size[0]*frame.size[1]

                if save_cropped == seqid:
                    filename = cropped_folder / f'{seqid}_{len(images):02}.jpg'
                    with open(filename, 'wb') as f:
                        frame.save(f)

                frame = to_tensor(frame)

                # added
                if args.single_channel:
                    frame = frame[2:3:, :, :]
                    frame = torch.cat((frame, frame, frame), dim = 0)
                    
                images.append(frame)
            # This creates a new dimension ('L'), L×N×H×W. Stacking is done on the first dimension (index 0),
            # creating a 4D tensor in which every frame is stacked along the new dimension.
            tensor = {'class': class_value, 'tensor': torch.stack(images)}
            # Save tensor to file
            filename = tensors_folder / f'{seqid}.lz4'
            with lz4.frame.open(filename, 'wb') as f:
                pickle.dump(tensor, f)

    make_tensors(train_ids, 'train')
    make_tensors(test_ids, 'test')

    mean = sum_x / num_pixels
    variance = sum_x2 / num_pixels - mean ** 2
    stddev = np.sqrt(variance)

    filename = Path(args.tensors_folder) / 'info.pkl'
    with open(filename, 'wb') as f:
        pickle.dump((class_names, mean, stddev), f)


class MyArgumentParser(argparse.ArgumentParser):
    def convert_arg_line_to_args(self, arg_line: str) -> list[str]:
        arg_line = arg_line.partition('#')[0].strip()  # Keep a list with only the first element of the partition
        if arg_line == '':  # Skip empty lines
            return []
        return arg_line.replace('=', ',').split(',')


if __name__ == '__main__':
    parser = MyArgumentParser(fromfile_prefix_chars='@')
    subparsers = parser.add_subparsers(required=True)

    # create the parser for the "make_sequences" command
    parser_mksq = subparsers.add_parser('make_sequences', help='Generate sequences file')
    parser_mksq.add_argument('source_folder', help='Folder containing original sequences divided into classes (and subfolders)')
    parser_mksq.add_argument('sequences_filename', help='Filename for the extracted sequences')
    parser_mksq.set_defaults(func=cmd_make_sequences)

    # create the parser for the "make_split" command
    parser_split = subparsers.add_parser('make_split', help='Generate splits')
    group_split = parser_split.add_mutually_exclusive_group(required=True)
    group_split.add_argument('--numtrain', type=int, help='Number of samples in training set (uniformly sampled from all classes)')
    group_split.add_argument('--perctrain', type=int, choices=range(1, 100), help='% of samples per class in training set (1-100)')
    parser_split.add_argument('sequences_filename', help='Filename for the extracted sequences')
    parser_split.add_argument('split_filename', help='Filename for the split to create')
    parser_split.set_defaults(func=cmd_make_split)

    # create the parser for the "make_tensors" command
    parser_mktn = subparsers.add_parser('make_tensors', help='Generate cropped tensors')
    parser_mktn.add_argument('--crop_box', type=int, nargs=4, help='Crop box values')
    parser_mktn.add_argument('--single_channel', type=bool, help='Remove red difference channel')
    parser_mktn.add_argument('--cropped_folder', help='Folder where the cropped images will be saved (empty=no save)')
    parser_mktn.add_argument('split_filename', help='Filename with the input split')
    parser_mktn.add_argument('tensors_folder', help='Folder to save the output tensors')
    parser_mktn.set_defaults(func=cmd_make_tensors)

  
    args = parser.parse_args()
    args.func(args)  # Call the appropriate function