tools/calibration/infer_raw_results.py

   1 """
   2 Copyright (C) 2018-2019 Intel Corporation
   3
   4 Licensed under the Apache License, Version 2.0 (the "License");
   5 you may not use this file except in compliance with the License.
   6 You may obtain a copy of the License at
   7
   8       http://www.apache.org/licenses/LICENSE-2.0
   9
  10 Unless required by applicable law or agreed to in writing, software
  11 distributed under the License is distributed on an "AS IS" BASIS,
  12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 See the License for the specific language governing permissions and
  14 limitations under the License.
  15 """
  16
  17 import os
  18 import numpy
  19 import pickle
  20 import shutil
  21 import tempfile
  22 from typing import Dict
  23
  24
  25 class InferRawResults:
  26     def __init__(self):
  27         self._size = 0
  28         self._index = 0
  29         self._dir_path = None
  30         pass
  31
  32     def release(self):
  33         if self._dir_path:
  34             shutil.rmtree(self._dir_path)
  35             self._dir_path = None
  36
  37     def __iter__(self):
  38         self._index = 0
  39         return self
  40
  41     def __next__(self):
  42         if self._index < self._size:
  43             file_path = os.path.join(self._dir_path, str(self._index))
  44             self._index += 1
  45
  46             f = open(file_path, "rb")
  47             try:
  48                 loaded_value = pickle.load(f)
  49             finally:
  50                 f.close()
  51             return loaded_value
  52         else:
  53             raise StopIteration
  54
  55     def size(self):
  56         return self._size
  57
  58     def add(self, value: Dict[str, numpy.ndarray]):
  59         if self._dir_path is None:
  60             self._dir_path = tempfile.mkdtemp("__infer_raw_results")
  61             if not os.path.exists(self._dir_path):
  62                 os.makedirs(self._dir_path)
  63
  64         file_path = os.path.join(self._dir_path, str(self._size))
  65
  66         f = open(file_path, "wb")
  67         try:
  68             pickle.dump(value, f)
  69         finally:
  70             f.close()
  71
  72         self._size += 1