ImageQualityMeasure.py 1.58 KB
Newer Older
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
1
import logging
2

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
3
import numpy as np
4 5
from bob.ip.qualitymeasure import galbally_iqm_features as iqm
from bob.ip.qualitymeasure import msu_iqa_features as iqa
6
from sklearn.preprocessing import FunctionTransformer
7 8 9

logger = logging.getLogger(__name__)

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
10

11 12 13
def iqm_features(images, galbally=True, msu=True, dtype=None):
    if not (galbally or msu):
        raise ValueError("At least galbally or msu needs to be True.")
14

15 16 17
    all_features = []
    for data in images:
        assert isinstance(data, np.ndarray)
18

19
        features = []
20

21
        if galbally:
22

23
            try:
24

25 26 27
                gf_set = iqm.compute_quality_features(data)
                gf_set = np.nan_to_num(gf_set)
                features = np.hstack((features, gf_set))
28

29
            except Exception:
30

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
31
                logger.error("Failed to extract galbally features.", exc_info=True)
32

33
                features = np.zeros((18,))
34

35
        if msu:
36

37
            try:
38

39 40 41
                msuf_set = iqa.compute_msu_iqa_features(data)
                msuf_set = np.nan_to_num(msuf_set)
                features = np.hstack((features, msuf_set))
42

43
            except Exception:
44

45
                logger.error("Failed to extract MSU features.", exc_info=True)
46

47
                features = np.zeros((121,))
48

49
        if dtype is not None:
50

51 52
            features = features.astype(dtype)
        all_features.append(features)
53

54
    return np.array(all_features)
55

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
56

57 58 59
def ImageQualityMeasure(galbally=True, msu=True, dtype=None, **kwargs):
    kw_args = dict(galbally=galbally, msu=msu, dtype=dtype)
    return FunctionTransformer(iqm_features, validate=False, kw_args=kw_args)