require 'torch';

require 'image';

folder = '/home/andrew/mitosis-detection'

train_folder = paths.concat(folder,'mitosis-train')

test_folder = paths.concat(folder,'mitosis-test')

trainset = paths.concat(folder,'mitosis-train.t7')

testset = paths.concat(folder,'mitosis-test.t7')

function gen_dataset(data_folder, outfile)

    true_folder = paths.concat(data_folder,'true')

    false_folder = paths.concat(data_folder,'false')

    true_files = {}

    false_files = {}

    for file in paths.files(true_folder,'.png') do

        table.insert(true_files, paths.concat(true_folder,file))

    end

    for file in paths.files(false_folder,'.png') do

        table.insert(false_files, paths.concat(false_folder,file))

    end

    data = torch.ByteTensor(#true_files+#false_files,3,101,101)

    label = torch.ByteTensor(#true_files+#false_files)

    for i,file in ipairs(true_files) do

        print(i)

        data[{{i},{},{},{}}] = image.load(file,3,'byte')

        label[i] = 1

    end

    for i,file in ipairs(false_files) do

        print(i + #true_files)

        data[{{i + #true_files},{},{},{}}] = image.load(file,3,'byte')

        label[i + #true_files] = 1

    end

    dataset = {}

    dataset.data = data

    dataset.label = label

    torch.save(outfile, dataset)

end

gen_dataset(train_folder,trainset)

gen_dataset(test_folder,testset)