<gh_stars>1-10
import yaml
from models import VGGNet
import pandas as pd
import os
from random import shuffle
from glob import glob

        
def get_model(input_shape, setting):
    setting['nc'] = len(setting['classes'])
    network = VGGNet(name='vgg16', ch=input_shape[-1], imgsz=input_shape, setting=setting)

    return network


def parse_dataset(path, name='trainval'):
    if name == 'trainval':
        imagepath = path['images']
        labelpath = path['labels']

        label_df = pd.read_csv(labelpath)
        # imagelist = glob(f"{imagepath}/**/**/*.dcm", recursive=True)
        rids = os.listdir(imagepath)
        shuffle(rids)
        valid_rate = 0.1

        train_rids = rids[:int(len(rids) * (1 - valid_rate))]
        valid_rids = rids[int(len(rids) * (1 - valid_rate)):]

        train_imagelist, valid_imagelist = [], []

        for rid in train_rids:
            templist = glob(f"{imagepath}/{rid}/**/*.dcm", recursive=True)
            train_imagelist += templist

        for rid in valid_rids:
            templist = glob(f"{imagepath}/{rid}/**/*.dcm", recursive=True)
            valid_imagelist += templist

        train_label_df = label_df[label_df["BraTS21ID"].isin(list(map(int, train_rids)))]
        valid_label_df = label_df[label_df["BraTS21ID"].isin(list(map(int, valid_rids)))]

        return [{"images": train_imagelist, "labels": train_label_df, "rids": train_rids}, 
                {"images": valid_imagelist, "labels": valid_label_df, "rids": valid_rids}]
    
    elif name == 'test':
        imagepath = path['images']

        imagelist = glob(f"{imagepath}/**/**/*.dcm", recursive=True)
        return {"images": imagelist}


if __name__ == "__main__":
    with open(f'./setting.yaml') as f:
        setting = yaml.load(f, Loader=yaml.SafeLoader)
    print(setting)

    model = get_model(input_shape=(64, 64, 1), setting=setting)
    print(model)

    train_path = setting['train']
    test_path = setting['test']

    train_dict = parse_dataset(train_path, name='trainval')
    test_dict = parse_dataset(test_path, name='test')

    model.fit(train_data=train_dict[0], validation_data=train_dict[0], epochs=100, batch_size=1)
import yaml
from models import VGGNet
import pandas as pd
import os
from random import shuffle
from glob import glob

        
def get_model(input_shape, setting):
    setting['nc'] = len(setting['classes'])
    network = VGGNet(name='vgg16', ch=input_shape[-1], imgsz=input_shape, setting=setting)

    return network


def parse_dataset(path, name='trainval'):
    if name == 'trainval':
        imagepath = path['images']
        labelpath = path['labels']

        label_df = pd.read_csv(labelpath)
        # imagelist = glob(f"{imagepath}/**/**/*.dcm", recursive=True)
        rids = os.listdir(imagepath)
        shuffle(rids)
        valid_rate = 0.1

        train_rids = rids[:int(len(rids) * (1 - valid_rate))]
        valid_rids = rids[int(len(rids) * (1 - valid_rate)):]

        train_imagelist, valid_imagelist = [], []

        for rid in train_rids:
            templist = glob(f"{imagepath}/{rid}/**/*.dcm", recursive=True)
            train_imagelist += templist

        for rid in valid_rids:
            templist = glob(f"{imagepath}/{rid}/**/*.dcm", recursive=True)
            valid_imagelist += templist

        train_label_df = label_df[label_df["BraTS21ID"].isin(list(map(int, train_rids)))]
        valid_label_df = label_df[label_df["BraTS21ID"].isin(list(map(int, valid_rids)))]

        return [{"images": train_imagelist, "labels": train_label_df, "rids": train_rids}, 
                {"images": valid_imagelist, "labels": valid_label_df, "rids": valid_rids}]
    
    elif name == 'test':
        imagepath = path['images']

        imagelist = glob(f"{imagepath}/**/**/*.dcm", recursive=True)
        return {"images": imagelist}


if __name__ == "__main__":
    with open(f'./setting.yaml') as f:
        setting = yaml.load(f, Loader=yaml.SafeLoader)
    print(setting)

    model = get_model(input_shape=(64, 64, 1), setting=setting)
    print(model)

    train_path = setting['train']
    test_path = setting['test']

    train_dict = parse_dataset(train_path, name='trainval')
    test_dict = parse_dataset(test_path, name='test')

    model.fit(train_data=train_dict[0], validation_data=train_dict[0], epochs=100, batch_size=1)
# imagelist = glob(f"{imagepath}/**/**/*.dcm", recursive=True)