all_calc_dataset.py

'''
2024.6.19: calculate the class distribution of training dataset
'''
import cv2
import os
import argparse
from dataset.vaihingen_dataset_cv2 import VaihingenDataSet
import numpy as np
import pandas as pd

DATA_PATH_HOME = '/localssd/yinxia/data'
SAVE_WEIGHT_PATH = '/localssd/yinxia/CRGNetdata/'
IGNORE_VALUE = 255

def get_args(city="Vaihingen"):
    parser = argparse.ArgumentParser()
    parser.add_argument("--city", default=city)
    args = parser.parse_args()
    # dataset
    if args.city=="Vaihingen":
        args.data_dir_train = f'{DATA_PATH_HOME}/Vaihingen/'
        args.data_dir_val = args.data_dir_test = args.data_dir_train
        args.train_list = f'./dataset/{args.city.lower()}_train.txt'
        args.test_list = args.val_list =f'./dataset/{args.city.lower()}_test.txt'
        args.name_classes = ['impervious surfaces', 'buildings', 'low vegetation', 'trees', 'cars']
        # BGR
        args.palette = [(255, 255, 255), (255, 0, 0), (255, 255, 0), (0, 255, 0), (0, 255, 255), (0, 0, 255), (0, 0, 0)]

    elif args.city=="Zurich":
        args.data_dir_train = f'{DATA_PATH_HOME}/Zurich/'
        args.data_dir_val = args.data_dir_test = args.data_dir_train
        args.train_list = f'./dataset/{args.city.lower()}_train.txt'
        args.test_list = args.val_list = f'./dataset/{args.city.lower()}_test.txt'
        args.name_classes = ['Roads', 'Buildings', 'Trees', 'Grass', 'Bare Soil', 'Water', 'Rails', 'Pools']
        args.palette = [(0, 0, 0), (100, 100, 100), (0, 125, 0), (0, 255, 0), (0, 80, 150),
                        (150, 0, 0), (0, 255, 255), (255, 150, 150), (255, 255, 255)]

    elif args.city=="uavid":
        args.data_dir_train = f'{DATA_PATH_HOME}/uavid/uavid_train/' # for training
        args.data_dir_val = f'{DATA_PATH_HOME}/uavid/uavid_val/' # for validation: select 7 images, each for each video
        args.data_dir_test = f'{DATA_PATH_HOME}/uavid/uavid_test/'
        args.train_list = f'./dataset/{args.city.lower()}_train.txt'
        args.val_list = f'./dataset/{args.city.lower()}_val.txt'
        args.test_list = f'./dataset/{args.city.lower()}_test.txt'
        args.name_classes = ['Building', 'Road', 'Tree', 'LowVeg', 'Cars'] # only consider five classes
        args.input_size_train = '1024, 1024'
        args.input_size_test = '1024, 1024'
        args.batch_size = 8
        args.num_workers = 4
        args.cache = True # False
        args.palette = [(0, 0, 128), (128, 64, 128), (0, 128, 0), (0, 128, 128), (128, 0, 64), (192, 0, 192),
                        (0, 64, 64), (0, 0, 0)]
        # if 'TELT' in args.method:
        #     args.batch_size = 4
        #     args.num_workers = 4
    else:
        raise ValueError('unknown city')
    args.num_classes = len(args.name_classes)
    
    return args


def calclass(filelist, num_classes):
    num = len(filelist)
    stats = np.zeros((num, num_classes+1), dtype=np.uint64)
    namelist = []
    for idx, file in enumerate(filelist):
        namelist.append(file['name'])

        label = cv2.imread(file['label'], cv2.IMREAD_UNCHANGED).flatten()
        # the total values
        stats[idx, -1] = label.size

        label = label[label!=IGNORE_VALUE]

        unique, counts = np.unique(label, return_counts=True)
        stats[idx, unique] = counts

    return stats, namelist


if __name__=="__main__":
    
    savepath = './dataset/stats'
    os.makedirs(savepath, exist_ok=True)

    # city = 
    # city = 
    citylist = ["Vaihingen","Zurich", "uavid"]
    for city in citylist:
        args = get_args(city)
        num_classes = int(args.num_classes)

        # train with all labels  # train with point labels # test with all labels
        datasetlist = [VaihingenDataSet(args.data_dir_train, args.train_list, mode=0),
        VaihingenDataSet(args.data_dir_train, args.train_list, mode=1, id=1),
        VaihingenDataSet(args.data_dir_test, args.test_list, mode=0)]

        suffixlist = ['trainall', 'trainpoint', 'testall']
        # calculate each type 
        for idx, dataset in enumerate(datasetlist):
            filelist = dataset.files
            # print(filelist[0])
            # print(len(filelist))
            stats, namelist = calclass(filelist,num_classes)
            #
            savefile = os.path.join(savepath, f'{city}_{suffixlist[idx]}.csv')
            #
            df = pd.DataFrame(stats)
            df['name'] = namelist

            df.to_csv(savefile)