[96354c]: / src / dataset / utils / dataset.py

Download this file

27 lines (23 with data), 979 Bytes

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import os
from typing import List, Tuple
import numpy as np
import csv
from src.dataset.patient import Patient
def read_brats(csv_path: str, lgg_only: bool=False) -> Tuple[List, List]:
patients_test = []
patients_train = []
with open(csv_path, 'r') as csvfile:
reader = csv.reader(csvfile, skipinitialspace=True)
next(reader, None)
for row in reader:
if lgg_only and row[1] == "HGG":
continue
patients_train.append(Patient(idx=row[0], center=row[3], grade=row[1], patient=row[2], patch_name=row[4],
size=list(map(int, row[5].split("x"))), data_path=os.path.dirname(csv_path),
train=True))
return patients_train, patients_test
def create_roi_mask(data: np.ndarray) -> np.ndarray:
# filter values bigger than 0
brain_mask = np.zeros(data.shape, np.float)
brain_mask[data > 0] = 1
return brain_mask