Skip to content
GitLab
Switch branch/tag
Find file Normal viewHistoryPermalink
_data_structures.py 5.32 KB
Newer Older
Matteo's avatar
update  
Matteo committed
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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129 
from dataclasses import dataclass
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.neighbors import KNeighborsClassifier

from mpai_cae_arp.audio.standards import EqualizationStandard, SpeedStandard
from ._constants import INVERSE_CLASSIFICATION_MAPPING


@dataclass
class ClassificationResult:
    """
    A class to represent the result of a classification.

    Since the classification recognizes 4 informations, those are stored in a container class to have the possibility to find all informations in the same place, but, when necessary, use only the part that is needed.

    Informations are:

    - the reading speed of the tape,
    - the writing speed,
    - the reading post-emphasis equalization curve
    - the writing pre-emphasis equalization curve
    """

    writing_speed: SpeedStandard
    reading_speed: SpeedStandard
    writing_equalization: EqualizationStandard
    reading_equalization: EqualizationStandard


_MAP_CLASS_TO_RESULT = {
    '3N_3N':
    ClassificationResult(SpeedStandard.III, SpeedStandard.III,
                         EqualizationStandard.NAB, EqualizationStandard.NAB),
    '3N_7C':
    ClassificationResult(SpeedStandard.III, SpeedStandard.IV,
                         EqualizationStandard.NAB, EqualizationStandard.CCIR),
    '3N_7N':
    ClassificationResult(SpeedStandard.III, SpeedStandard.IV,
                         EqualizationStandard.NAB, EqualizationStandard.NAB),
    '3N_15C':
    ClassificationResult(SpeedStandard.III, SpeedStandard.V,
                         EqualizationStandard.NAB, EqualizationStandard.CCIR),
    '3N_15N':
    ClassificationResult(SpeedStandard.III, SpeedStandard.V,
                         EqualizationStandard.NAB, EqualizationStandard.NAB),
    '7C_3N':
    ClassificationResult(SpeedStandard.IV, SpeedStandard.III,
                         EqualizationStandard.CCIR, EqualizationStandard.NAB),
    '7C_7C':
    ClassificationResult(SpeedStandard.IV, SpeedStandard.IV,
                         EqualizationStandard.CCIR, EqualizationStandard.CCIR),
    '7C_7N':
    ClassificationResult(SpeedStandard.IV, SpeedStandard.IV,
                         EqualizationStandard.CCIR, EqualizationStandard.NAB),
    '7C_15C':
    ClassificationResult(SpeedStandard.IV, SpeedStandard.V,
                         EqualizationStandard.CCIR, EqualizationStandard.CCIR),
    '7C_15N':
    ClassificationResult(SpeedStandard.IV, SpeedStandard.V,
                         EqualizationStandard.CCIR, EqualizationStandard.NAB),
    '7N_3N':
    ClassificationResult(SpeedStandard.IV, SpeedStandard.III,
                         EqualizationStandard.NAB, EqualizationStandard.NAB),
    '7N_7C':
    ClassificationResult(SpeedStandard.IV, SpeedStandard.IV,
                         EqualizationStandard.NAB, EqualizationStandard.CCIR),
    '7N_7N':
    ClassificationResult(SpeedStandard.IV, SpeedStandard.IV,
                         EqualizationStandard.NAB, EqualizationStandard.NAB),
    '7N_15C':
    ClassificationResult(SpeedStandard.IV, SpeedStandard.V,
                         EqualizationStandard.NAB, EqualizationStandard.CCIR),
    '7N_15N':
    ClassificationResult(SpeedStandard.IV, SpeedStandard.V,
                         EqualizationStandard.NAB, EqualizationStandard.NAB),
    '15C_3N':
    ClassificationResult(SpeedStandard.V, SpeedStandard.III,
                         EqualizationStandard.CCIR, EqualizationStandard.NAB),
    '15C_7C':
    ClassificationResult(SpeedStandard.V, SpeedStandard.IV,
                         EqualizationStandard.CCIR, EqualizationStandard.CCIR),
    '15C_7N':
    ClassificationResult(SpeedStandard.V, SpeedStandard.IV,
                         EqualizationStandard.CCIR, EqualizationStandard.NAB),
    '15C_15C':
    ClassificationResult(SpeedStandard.V, SpeedStandard.V,
                         EqualizationStandard.CCIR, EqualizationStandard.CCIR),
    '15C_15N':
    ClassificationResult(SpeedStandard.V, SpeedStandard.V,
                         EqualizationStandard.CCIR, EqualizationStandard.NAB),
    '15N_3N':
    ClassificationResult(SpeedStandard.V, SpeedStandard.III,
                         EqualizationStandard.NAB, EqualizationStandard.NAB),
    '15N_7C':
    ClassificationResult(SpeedStandard.V, SpeedStandard.IV,
                         EqualizationStandard.NAB, EqualizationStandard.CCIR),
    '15N_7N':
    ClassificationResult(SpeedStandard.V, SpeedStandard.IV,
                         EqualizationStandard.NAB, EqualizationStandard.NAB),
    '15N_15C':
    ClassificationResult(SpeedStandard.V, SpeedStandard.V,
                         EqualizationStandard.NAB, EqualizationStandard.CCIR),
    '15N_15N':
    ClassificationResult(SpeedStandard.V, SpeedStandard.V,
                         EqualizationStandard.NAB, EqualizationStandard.NAB),
}


class Classifier:

    model: RandomForestClassifier | DecisionTreeClassifier | KNeighborsClassifier

    def __init__(self, model):
        self.model = model

    def predict(self, x: pd.DataFrame) -> pd.DataFrame:
        prediction = self.model.predict(x)
        prediction = pd.DataFrame(prediction, columns=['classification'])

        prediction = prediction.replace(INVERSE_CLASSIFICATION_MAPPING)

        prediction = prediction.replace(_MAP_CLASS_TO_RESULT)

        return prediction

    def get_model_description(self) -> str:
        return str(self.model)