methods/classes/LandCodeDataset.py - master - maratgaliulin/landcode_classifier

methods/classes/LandCodeDataset.py
81 строка3 KB
maratgaliulin
first commit
26 янв 2026, 02:03
26 янв 2026, 02:0318d5031
import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset
from sklearn.preprocessing import LabelEncoder

class LandCodeDataset(Dataset):
    def __init__(self, init_df:pd.DataFrame, label_col:str):
        """
        Обработка датасета
        """
        

        all_columns = init_df.columns

        feature_columns = [
            col for col in all_columns if col != label_col
        ]

        label_df = init_df[label_col]
        features_df = init_df[feature_columns]


        # print(f"DEBUG: label_df: {label_df}")
        print(f"DEBUG: форма features_df: {features_df.shape}")
        print(f"DEBUG: форма label_df : {label_df.shape}")
        print(f"DEBUG: тип label_df : {type(label_df)}")
        
        if hasattr(label_df, 'values'):
            print(f"DEBUG: форма label_df.values : {label_df.values.shape}")
        
        
        self.tfidf_features = features_df.values.astype(np.float32)
        
        
        if hasattr(label_df, 'values'):
            labels_raw = label_df.values
            print(f"DEBUG: форма labels_raw : {labels_raw.shape}")
            print(f"DEBUG: labels_raw[:5]: {labels_raw[:5] if len(labels_raw) > 5 else labels_raw}")
            
            
            if len(labels_raw.shape) > 1:
                self.labels_original = labels_raw.flatten()
                print(f"DEBUG: уплощенный датасет: {self.labels_original.shape}")
            else:
                self.labels_original = labels_raw
        else:
            self.labels_original = np.array(label_df).flatten()
        
        
        print(f"DEBUG: длина массива фичей: {len(self.tfidf_features)}")
        print(f"DEBUG: длина У-столбца: {len(self.labels_original)}")
        
        if len(self.tfidf_features) != len(self.labels_original):
            raise ValueError(f"Несоответствие длин массивов фичей и У-столбца!!! Фичи: {len(self.tfidf_features)}, У-столбец: {len(self.labels_original)}")
        
        
        self.label_encoder = LabelEncoder()
        self.labels_encoded = self.label_encoder.fit_transform(self.labels_original)
        self.num_classes = len(self.label_encoder.classes_)
        
        
        print(f"  форма датасета с фичами TF-IDF : {self.tfidf_features.shape}")
        print(f"  кол-во сэмплов: {len(self.tfidf_features)}")
        print(f"  кол-во уник.классов: {self.num_classes}")
        
    def __len__(self):
        return len(self.tfidf_features)
    
    def __getitem__(self, idx):

        # ******** DEBUG ***********:

        if idx >= len(self.labels_encoded):
            raise IndexError(f"Индекс {idx} за пределами списка. Длина labels_encoded : {len(self.labels_encoded)}")
        
        return {
            'tfidf_features': torch.FloatTensor(self.tfidf_features[idx]),
            'label': torch.LongTensor([self.labels_encoded[idx]]).squeeze(),
            'idx': idx
        }