Coverage for src/code_beatrix/ai/dlbase.py : 75%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2@file
3@brief Base class for deep learning models.
4"""
5import os
8class DeepLearningBase:
9 """
10 Implements a common interface to manipulate pre-trained
11 deep learning models.
12 """
14 def __init__(self, model, gpu, fLOG=None):
15 """
16 @param model model (url, filename, ...)
17 @param gpu use gpu
18 @param fLOG logging function
19 """
20 self._gpu = gpu
21 if model is None:
22 raise ValueError("model must be specified")
23 if isinstance(model, str):
24 if not os.path.exists(model):
25 raise FileNotFoundError("Unable to find '{0}'".format(model))
26 raise NotImplementedError(
27 "Unable to load model '{0}'".format(model))
28 self._model = model
29 self._fLOG = fLOG
31 def log(self, *args, **kwargs):
32 """
33 Log something.
34 """
35 if self._fLOG:
36 self._fLOG(*args, **kwargs)
38 def predict(self, X):
39 """
40 Applies the model on features *X*.
42 @param X features
43 @return prediction
44 """
45 raise NotImplementedError("Method predict is not implemented.")
48class DeepLearningImage(DeepLearningBase):
49 """
50 Implements a common interface to manipulate pre-trained
51 deep learning models processing images.
52 """
54 def __init__(self, model, gpu=False, fLOG=None):
55 """
56 @param model model name
57 @param gpu use gpu
58 @param fLOG logging function
59 """
60 DeepLearningBase.__init__(self, model, gpu=gpu, fLOG=fLOG)