dataset_preparation.py

from torch.utils.data import Dataset
from PIL import Image
import os
from typing import Optional, Tuple
import pandas as pd
import torch


class MalwareDetect2Dataset(Dataset):
    def __init__(self, data_dir: str, transform: Optional[callable] = None):
        self.data_dir = data_dir
        self.transform = transform
        self.data = self._load_data(data_dir)
        self.label_to_class = self._get_label_to_class_mapping(data_dir)
    
    def _load_data(self, data_dir: str) -> pd.DataFrame:
        data = []
        for label, sub_dir in enumerate(sorted(os.listdir(data_dir))):
            sub_dir_path = os.path.join(data_dir, sub_dir)
            if os.path.isdir(sub_dir_path):
                for img_file in os.listdir(sub_dir_path):
                    data.append({
                        'image_path': os.path.join(sub_dir_path, img_file),
                        'label': label
                    })
        return pd.DataFrame(data)

    def _get_label_to_class_mapping(self, data_dir: str) -> dict:
        return {label: class_name for label, class_name in enumerate(sorted(os.listdir(data_dir)))}
    
    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, idx: int) -> Tuple[Image.Image, int]:
        row = self.data.iloc[idx]
        image = Image.open(row['image_path'])
        label = row['label']

        if self.transform:
            image = self.transform(image)

        return image, label
    
    def class_counts(self, mapping: Optional[bool]) -> dict:
        counts = self.data['label'].value_counts().to_dict()
        if mapping:
            return {self.label_to_class[label]: count for label, count in counts.items()}
        return counts
    
class OpcodeDataset(Dataset):
    def __init__(self, features, labels):
        """
        Args:
        - features (numpy.ndarray): Feature matrix (e.g., loaded from X_train.npy).
        - labels (numpy.ndarray): Corresponding labels (e.g., loaded from y_train.npy).
        """
        self.features = torch.tensor(features, dtype=torch.float32)
        self.labels = torch.tensor(labels, dtype=torch.long)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        """
        Returns a single sample and its label at the specified index.
        """
        return self.features[idx], self.labels[idx]