PhotometryClassificationWebservice/random_forest_classifier.py at master · BHTOM-Team/PhotometryClassificationWebservice · GitHub

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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import logging
import pickle
from typing import List

import numpy as np
import pandas as pd
from sklearn import metrics
from sklearn.ensemble import RandomForestClassifier

from sklearn.model_selection import KFold

from classification_source import ClassificationSource, SELECTED_FEATURES, TRAINING_FILE_PATHS, MODEL_PATHS


logger: logging.Logger = logging.getLogger()
logging.basicConfig(level=logging.DEBUG)


def empty_training_file(source: ClassificationSource) -> pd.DataFrame:
    return pd.DataFrame(columns=[[*SELECTED_FEATURES.get(source, []), 'class']])


def read_training_data(source: ClassificationSource) -> pd.DataFrame:
    filename: str = TRAINING_FILE_PATHS.get(source, '')
    try:
        gaia_df: pd.DataFrame = pd.read_csv(filename)[[*SELECTED_FEATURES.get(source, []), 'class']]
        return gaia_df
    except FileNotFoundError:
        logger.error(f'File {filename} not found!')
        return empty_training_file(source)
    except pd.errors.ParserError as e:
        logger.error(f'Parser error while reading {filename}: {e}')
        return empty_training_file(source)
    except Exception as e:
        logger.error(f'Unexpected error while reading {filename}: {e}')
        return empty_training_file(source)


def train_rfc(source: ClassificationSource) -> RandomForestClassifier:
    logger.info(f'Training RFC for {source.name}...')
    np.random.seed(42)
    rfc: RandomForestClassifier = RandomForestClassifier(n_estimators=500, random_state=42)
    kf = KFold(n_splits=9, random_state=42, shuffle=True)
    training_data: pd.DataFrame = read_training_data(source)
    selected_features: List[str] = SELECTED_FEATURES.get(source, [])
    kf.get_n_splits(training_data)
    for train_index, test_index in kf.split(training_data):
        train, test = training_data.loc[train_index], training_data.loc[test_index]
        train_features, train_class = train.loc[:, [*selected_features]], train.loc[:, 'class']

        test_features, test_class = test.loc[:, [*selected_features]], test.loc[:, 'class']

        rfc.fit(train_features, train_class)

        rfc_sc = metrics.accuracy_score(test_class, rfc.predict(test_features))
        print(f'RFC model accuracy: {rfc_sc}, for {source.name}')
    return rfc


def load_rfc(source: ClassificationSource) -> RandomForestClassifier:
    try:
        logger.info(f'Loading RFC for {source.name} from file...')
        with open(MODEL_PATHS.get(source, ''), 'rb') as handle:
            rfc: RandomForestClassifier = pickle.load(handle)
            return rfc
    except Exception as e:
        logger.error(f'Exception {e} when loading RFC for {source.name} from file!')
        rfc: RandomForestClassifier = train_rfc(source)
        with open(MODEL_PATHS.get(source, ''), 'wb') as handle:
            logger.info(f'Saving RFC for {source.name} to file...')
            pickle.dump(rfc, handle)
        return rfc