<reponame>bkahn-github/track-ml
import os
import logging

import numpy as np

import torch
from torch.utils.data import Dataset, DataLoader

import torchvision
from torchvision import transforms, utils

from trackml.dataset import load_event

def load_data(root_path, sub_folder, event, subset):
    logging.info('-' * 50)
    logging.info('Loading data')

    total_path = os.path.join(root_path, sub_folder, event)

    hits, cells, particles, truth = load_event(total_path)

    if subset != -1: 
        logging.info('Returning a subset of the data')

        hits = hits.sample(subset)
        truth = truth[truth.hit_id.isin(hits.hit_id.values)]
        particles = particles[particles.particle_id.isin(truth.particle_id.values)]

    return hits, cells, particles, truth

class TripleNetDataset(Dataset):
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        data = self.data.iloc[[idx]]

        x = data[['hit_id-1', 'hit_id-2', 'hit_id-3', 'rho-1', 'phi-1', 'rho-2', 'phi-2', 'rho-3', 'phi-3']].values[0]
        y = data['true_track'].values[0]
        y = np.array(y)

        x = torch.from_numpy(x).float()
        y = torch.from_numpy(y).float()

        return {'x': x, 'y': y}
import os
import logging

import numpy as np

import torch
from torch.utils.data import Dataset, DataLoader

import torchvision
from torchvision import transforms, utils

from trackml.dataset import load_event

def load_data(root_path, sub_folder, event, subset):
    logging.info('-' * 50)
    logging.info('Loading data')

    total_path = os.path.join(root_path, sub_folder, event)

    hits, cells, particles, truth = load_event(total_path)

    if subset != -1: 
        logging.info('Returning a subset of the data')

        hits = hits.sample(subset)
        truth = truth[truth.hit_id.isin(hits.hit_id.values)]
        particles = particles[particles.particle_id.isin(truth.particle_id.values)]

    return hits, cells, particles, truth

class TripleNetDataset(Dataset):
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        data = self.data.iloc[[idx]]

        x = data[['hit_id-1', 'hit_id-2', 'hit_id-3', 'rho-1', 'phi-1', 'rho-2', 'phi-2', 'rho-3', 'phi-3']].values[0]
        y = data['true_track'].values[0]
        y = np.array(y)

        x = torch.from_numpy(x).float()
        y = torch.from_numpy(y).float()

        return {'x': x, 'y': y}