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

methods/classes/LandUsePredictor.py
429 строк19 KB
maratgaliulin
21 июн 2026, 18:12
21 июн 2026, 18:122f0edd6
import torch
import pickle
import os
import json
import re
import pandas as pd
import numpy as np

from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report
import torch.nn.functional as F


from openpyxl.utils import get_column_letter
from openpyxl.styles import PatternFill
from openpyxl.formatting.rule import FormulaRule
from torch.utils.data import DataLoader, TensorDataset
from transformers import AutoModel
from tqdm.auto import tqdm
from .BertWithNumeric import BertWithNumeric
from ..utils import return_is_in_text_columns

class LandUsePredictor:
    def __init__(self, model_path, components_path):

        with open(components_path, 'rb') as f:
            components = pickle.load(f)
        
        self.tokenizer = components['tokenizer']
        self.label_encoder = components['label_encoder']
        self.text_columns = None
        
        # Load model config
        checkpoint = torch.load(model_path, map_location='cpu')
        self.text_columns = checkpoint['text_columns']
        config = checkpoint['model_config']
        
        bert_model = AutoModel.from_pretrained(config['bert_model_name'])

        classifier_in = checkpoint['model_state_dict']['classifier.weight'].shape[1]
        use_numeric_area = classifier_in == bert_model.config.hidden_size + 1
        self.area_scaler = None
        if use_numeric_area:
            scaler_path = components_path.replace('.pkl', '_area_scaler.pkl')
            if os.path.isfile(scaler_path):
                with open(scaler_path, 'rb') as f:
                    self.area_scaler = pickle.load(f)

        self.model = BertWithNumeric(
            bert_model, num_labels=config['num_labels'], use_numeric_area=use_numeric_area,
        )
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.model.eval()

        files_dir = os.path.abspath(os.path.join(os.path.dirname(__file__),  '../../files/'))
        
        self.priority1_mapping = self._load_mapping(os.path.join(files_dir, 'mapping_priority1.json'))
        self.priority2_mapping = self._load_mapping(os.path.join(files_dir, 'mapping_priority2.json'))
        
        
        # Set device
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model.to(self.device)

    @staticmethod
    def _load_mapping(path):
        if os.path.isfile(path):
            try:
                with open(path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                return {str(k): str(v) for k, v in data.items()}
            except Exception:
                return {}
        return {}

    @staticmethod
    def _normalize_text(value):
        if pd.isna(value):
            return ''
        s = str(value).lower()
        s = s.replace('ё', 'е')
        s = re.sub(r'\s+', ' ', s).strip()
        return s
    
    def preprocess_dataframe(self, df):        
        df_processed = df.copy()
        
        if 'combined_text' not in df_processed.columns:
            df_processed['combined_text'] = df_processed[self.text_columns].fillna('').astype(str).agg(' | '.join, axis=1)
        
        return df_processed
    
    def predict(self, df, batch_size=32, return_proba=False, show_progress=True):
        """Make predictions on new data with progress bar"""
        df_processed = self.preprocess_dataframe(df)
        
        # Tokenize
        encodings = self.tokenizer(
            df_processed['combined_text'].tolist(),
            truncation=True,
            padding=True,
            max_length=256,
            return_tensors='pt'
        )

        if self.model.use_numeric_area:
            if 'Площадь, кв.м' not in df_processed.columns:
                df_processed['Площадь, кв.м'] = 0.0
            areas = df_processed['Площадь, кв.м'].fillna(0).astype(float).values.reshape(-1, 1)
            if self.area_scaler is not None:
                areas = self.area_scaler.transform(areas).flatten()
            else:
                areas = areas.flatten()
            dataset = TensorDataset(
                encodings['input_ids'],
                encodings['attention_mask'],
                torch.tensor(areas, dtype=torch.float32),
            )
        else:
            dataset = TensorDataset(
                encodings['input_ids'],
                encodings['attention_mask'],
            )
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
        
        all_predictions = []
        all_probabilities = []
        
        # Create progress bar
        if show_progress:
            pbar = tqdm(total=len(dataloader), 
                        desc="Predicting", 
                        unit="batch",
                        bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]")
        
        with torch.no_grad():
            for batch_idx, batch in enumerate(dataloader):
                if self.model.use_numeric_area:
                    input_ids, attention_mask, area = [b.to(self.device) for b in batch]
                    outputs = self.model(input_ids, attention_mask, area)
                else:
                    input_ids, attention_mask = [b.to(self.device) for b in batch]
                    outputs = self.model(input_ids, attention_mask)
                
                if return_proba:
                    probabilities = torch.softmax(outputs, dim=1)
                    all_probabilities.append(probabilities.cpu().numpy())
                
                predictions = torch.argmax(outputs, dim=1)
                all_predictions.append(predictions.cpu().numpy())
                
                # Update progress bar
                if show_progress:
                    pbar.update(1)
                    pbar.set_postfix({
                        'samples': (batch_idx + 1) * batch_size,
                        'current_batch': batch_idx + 1
                    })
        
        # Close progress bar
        if show_progress:
            pbar.close()
        
        # Rest of the method remains the same...
        all_predictions = np.concatenate(all_predictions)
        decoded_labels = self.label_encoder.inverse_transform(all_predictions)
        
        result_df = df.copy()
        result_df['predicted_code'] = decoded_labels
        
        if return_proba:
            all_probabilities = np.concatenate(all_probabilities, axis=0)
            result_df['prediction_confidence'] = np.max(all_probabilities, axis=1)
            
            # Add all probability columns efficiently
            prob_data = {
                f'prob_{class_name}': all_probabilities[:, i]
                for i, class_name in enumerate(self.label_encoder.classes_)
            }
            prob_df = pd.DataFrame(prob_data, index=result_df.index)
            result_df = pd.concat([result_df, prob_df], axis=1)
        
        return result_df
    
    def predict_single(self, row_dict):
        df_single = pd.DataFrame([row_dict])
        result = self.predict(df_single, return_proba=True)
        
        return {
            'predicted_code': result['predicted_code'].iloc[0],
            'confidence': result['prediction_confidence'].iloc[0],
            'top_3': self.get_top_predictions(result, 3)
        }
    
    def get_top_predictions(self, result_df, top_n=3):
        prob_cols_exist = any(result_df.columns.str.startswith('prob_'))
    
        if not prob_cols_exist:
            result_df = self.predict(result_df, return_proba=True)
        
        prob_cols = [col for col in result_df.columns if col.startswith('prob_')]
        
        top_predictions = []
        for idx, row in result_df.iterrows():
            probs = [(col.replace('prob_', ''), row[col]) for col in prob_cols]
            probs_sorted = sorted(probs, key=lambda x: x[1], reverse=True)[:top_n]
            top_predictions.append(probs_sorted)
        
        return top_predictions
    
    def get_top_k(self, df, k=5):
        result_with_probs = self.predict(df, return_proba=True)
        prob_cols = [col for col in result_with_probs.columns if col.startswith('prob_')]
        
        top_k_results = []
        for idx, row in result_with_probs.iterrows():
            
            class_probs = []
            for col in prob_cols:
                class_name = col.replace('prob_', '')
                prob = row[col]
                class_probs.append((class_name, prob))
            
            class_probs.sort(key=lambda x: x[1], reverse=True)
            top_k = class_probs[:k]
            
            top_k_results.append({
                'predicted': row['predicted_code'],
                'confidence': row['prediction_confidence'],
                'top_k': top_k
            })
        
        return top_k_results
    
    def predict_data_from_client_csv_file(self, path_to_client_csv_file:str, sheet_name:str, path_to_saved_file:str) -> None:
        if(os.path.isfile(path_to_client_csv_file)):
            
            df = pd.read_excel(path_to_client_csv_file, sheet_name=sheet_name)
            dataframe_columns = df.columns.to_list()

            is_in_text_columns = return_is_in_text_columns(dataframe_columns=dataframe_columns, text_columns=self.text_columns)

            if(is_in_text_columns):

                priority_1_col = 'Вид использования участка по документу (САМЫЙ ГЛАВНЫЙ АТРИБУТ - Приоритет 1)'
                priority_2_col = 'Разрешенное использование (текстовое описание) - Приоритет 2'
                default_code_for_missing = '14:000'

                prefilled_mask = pd.Series(False, index=df.index)
                prefilled_code = pd.Series(index=df.index, dtype=object)
                prefilled_conf = pd.Series(index=df.index, dtype='float64')

                print(f"Общая длина датасета: {len(df)}")

                if priority_1_col in df.columns and priority_2_col in df.columns:
                    print('Нормализация данных в столбцах')
                    norm_p1 = df[priority_1_col].map(self._normalize_text)
                    norm_p2 = df[priority_2_col].map(self._normalize_text)
                    missing_both = (norm_p1 == '') & (norm_p2 == '')
                    if missing_both.any():
                        print(f'Данные отсутствуют в ключевых столбцах: {len(prefilled_code.loc[missing_both])}')
                        prefilled_mask |= missing_both
                        prefilled_code.loc[missing_both] = default_code_for_missing
                        prefilled_conf.loc[missing_both] = 1.0

                    if self.priority1_mapping:
                        print('Есть JSON файл для мэппинга Приоритета 1. Готовим маску для автопредсказания')
                        p1_mask = (~prefilled_mask) & norm_p1.isin(self.priority1_mapping)
                        if p1_mask.any():
                            prefilled_mask |= p1_mask
                            prefilled_code.loc[p1_mask] = norm_p1.loc[p1_mask].map(self.priority1_mapping)
                            prefilled_conf.loc[p1_mask] = 1.0

                    if self.priority2_mapping:
                        print('Есть JSON файл для мэппинга Приоритета 2. Готовим маску для автопредсказания')
                        p2_mask = (~prefilled_mask) & norm_p2.isin(self.priority2_mapping)
                        if p2_mask.any():
                            prefilled_mask |= p2_mask
                            prefilled_code.loc[p2_mask] = norm_p2.loc[p2_mask].map(self.priority2_mapping)
                            prefilled_conf.loc[p2_mask] = 1.0

                df_to_predict = df.loc[~prefilled_mask]
                if not df_to_predict.empty:
                    print(f'Есть автоматически предсказанные файлы - {len(df_to_predict)}')
                    predicted_dataframe = self.predict(df=df_to_predict, return_proba=True)
                    prefilled_code.loc[df_to_predict.index] = predicted_dataframe['predicted_code']
                    prefilled_conf.loc[df_to_predict.index] = predicted_dataframe['prediction_confidence']

                df['Предсказанный Код расчёта вида использования (ГБУ)'] = prefilled_code
                df['Уверенность модели'] = prefilled_conf.round(5) * 100
                predicted_code_col = 'Предсказанный Код расчёта вида использования (ГБУ)'

                mapping_path = './files/kody_zu.json'
                # mapping_path = os.path.abspath(os.path.join(os.path.dirname(__file__),  '../../files/kody_zu.json'))
                mapping_key_col = 'Код расчета вида использования ключ'
                mapping_cols = [
                    'Код вида разрешенного использования (в соответствии с классификатором видов разрешенного использования земельных участков)',
                    'Условный сегмент',
                    'Сегмент'
                ]

                if os.path.isfile(mapping_path):
                    print('Делаем мэппинг по кодам ЗУ')
                    mapping_df = pd.read_json(mapping_path)
                    available_cols = [mapping_key_col] + [col for col in mapping_cols if col in mapping_df.columns]
                    mapping_df = mapping_df[available_cols].copy()
                    mapping_df[mapping_key_col] = mapping_df[mapping_key_col].astype(str)
                    df[predicted_code_col] = df[predicted_code_col].astype(str)
                    df = df.merge(
                        mapping_df,
                        how='left',
                        left_on=predicted_code_col,
                        right_on=mapping_key_col
                    ).drop(columns=[mapping_key_col])
                
                output_sheet_name = 'предсказания'
                confidence_col_name = 'Уверенность модели'

                with pd.ExcelWriter(path_to_saved_file, engine='openpyxl') as writer:
                    df.to_excel(writer, sheet_name=output_sheet_name, index=False)

                    worksheet = writer.sheets[output_sheet_name]

                    if not df.empty and confidence_col_name in df.columns:
                        max_row = worksheet.max_row
                        max_col = worksheet.max_column

                        confidence_col_idx = df.columns.get_loc(confidence_col_name) + 1
                        confidence_col_letter = get_column_letter(confidence_col_idx)

                        start_cell = worksheet.cell(row=2, column=1).coordinate
                        end_cell = worksheet.cell(row=max_row, column=max_col).coordinate
                        data_range = f"{start_cell}:{end_cell}"

                        low_threshold = 0.1 * 100
                        mid_threshold = 0.4 * 100

                        red_fill = PatternFill(start_color="FFC7CE", end_color="FFC7CE", fill_type="solid")
                        yellow_fill = PatternFill(start_color="FFEB9C", end_color="FFEB9C", fill_type="solid")

                        red_rule = FormulaRule(formula=[f"${confidence_col_letter}2<{low_threshold}"], fill=red_fill)
                        yellow_rule = FormulaRule(
                            formula=[f"AND(${confidence_col_letter}2>={low_threshold},${confidence_col_letter}2<{mid_threshold})"],
                            fill=yellow_fill
                        )

                        worksheet.conditional_formatting.add(data_range, red_rule)
                        worksheet.conditional_formatting.add(data_range, yellow_rule)

                    # Highlight predicted code column in light blue
                    if not df.empty and predicted_code_col in df.columns:
                        predicted_col_idx = df.columns.get_loc(predicted_code_col) + 1
                        predicted_col_letter = get_column_letter(predicted_col_idx)
                        blue_fill = PatternFill(start_color="D9E8FF", end_color="D9E8FF", fill_type="solid")
                        for row in range(2, worksheet.max_row + 1):
                            cell = worksheet[f"{predicted_col_letter}{row}"]
                            cell.fill = blue_fill

                print('Работа модели завершена успешно')

            else:
                print("Столбцы, необходимые для работы, отсутствуют в наборе данных.")

        else:
            print("Введён неверный адрес датафрейма.")

    def evaluate_model(self, texts, true_labels, batch_size=32):
        """
        Оценка модели на тестовых данных
        
        Parameters:
        -----------
        texts : list
            Список текстов
        true_labels : list или array
            Истинные метки (индексы)
        batch_size : int
            Размер батча для инференса
        
        Returns:
        --------
        dict: Словарь с метриками (accuracy, loss, precision, recall, f1, classification_report)
        """
        
        
        self.model.eval()
        
        # Создаем датасет и даталоадер
        dataset = torch.utils.data.TensorDataset(
            self.tokenizer(texts, truncation=True, padding=True, max_length=256, return_tensors='pt')['input_ids'],
            self.tokenizer(texts, truncation=True, padding=True, max_length=256, return_tensors='pt')['attention_mask']
        )
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
        
        all_predictions = []
        all_probabilities = []
        total_loss = 0
        criterion = torch.nn.CrossEntropyLoss()
        
        with torch.no_grad():
            for batch in dataloader:
                input_ids, attention_mask = [b.to(self.device) for b in batch]
                outputs = self.model(input_ids, attention_mask)
                
                # Вычисляем loss
                labels_tensor = torch.tensor(true_labels[len(all_predictions):len(all_predictions)+len(input_ids)], device=self.device)
                loss = criterion(outputs, labels_tensor)
                total_loss += loss.item()
                
                probabilities = torch.softmax(outputs, dim=1)
                predictions = torch.argmax(outputs, dim=1)
                
                all_probabilities.append(probabilities.cpu().numpy())
                all_predictions.extend(predictions.cpu().numpy())
        
        # Метрики
        avg_loss = total_loss / len(dataloader)
        accuracy = accuracy_score(true_labels, all_predictions)
        precision = precision_score(true_labels, all_predictions, average='weighted', zero_division=0)
        recall = recall_score(true_labels, all_predictions, average='weighted', zero_division=0)
        f1 = f1_score(true_labels, all_predictions, average='weighted', zero_division=0)
        
        return {
            'loss': avg_loss,
            'accuracy': accuracy,
            'precision': precision,
            'recall': recall,
            'f1_score': f1,
            'classification_report': classification_report(true_labels, all_predictions, zero_division=0)
        }