Spaces:
Sleeping
Sleeping
File size: 1,654 Bytes
f5288df |
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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import yaml
import os
import torch
import random
import numpy as np
class BaseConfig:
def __init__(self):
config_path = os.path.join(os.path.dirname(__file__), 'config.yml')
with open(config_path, 'r') as file:
self.config = yaml.safe_load(file)
self.setup_environment()
def setup_environment(self):
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
os.environ['CUDA_VISIBLE_DEVICES'] = self.config["gpu"]["visible_device"]
self.device = torch.device(self.config["gpu"]["device"]) #torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
torch.set_float32_matmul_precision("medium")
def custom_collate(self, batch):
"""Handles variable size of the scans and pads the sequence dimension."""
images = [item['image'] for item in batch]
labels = [item['label'] for item in batch]
max_len = self.config["data"]["collate"] # Single scan input
padded_images = []
for img in images:
pad_size = max_len - img.shape[0]
if pad_size > 0:
padding = torch.zeros((pad_size,) + img.shape[1:])
img_padded = torch.cat([img, padding], dim=0)
padded_images.append(img_padded)
else:
padded_images.append(img)
return {"image": torch.stack(padded_images, dim=0), "label": torch.stack(labels)}
def get_config(self):
return self.config |