google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
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
"""utility function to read midas outputs."""
17
# from midas/utils.py
18
import re
19
import numpy as np
20

21

22
def read_pfm(path):
23
  """Read pfm file.
24

25
  Args:
26
      path (str): path to file
27

28
  Returns:
29
      tuple: (data, scale)
30

31
  Raises:
32
      Exception: for invalid pfm file
33
  """
34
  with open(path, "rb") as file:
35
    color = None
36
    width = None
37
    height = None
38
    scale = None
39
    endian = None
40

41
    header = file.readline().rstrip()
42
    if header.decode("ascii") == "PF":
43
      color = True
44
    elif header.decode("ascii") == "Pf":
45
      color = False
46
    else:
47
      raise Exception("Not a PFM file: " + path)
48

49
    dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii"))
50
    if dim_match:
51
      width, height = list(map(int, dim_match.groups()))
52
    else:
53
      raise Exception("Malformed PFM header.")
54

55
    scale = float(file.readline().decode("ascii").rstrip())
56
    if scale < 0:
57
      # little-endian
58
      endian = "<"
59
      scale = -scale
60
    else:
61
      # big-endian
62
      endian = ">"
63

64
    data = np.fromfile(file, endian + "f")
65
    shape = (height, width, 3) if color else (height, width)
66

67
    data = np.reshape(data, shape)
68
    data = np.flipud(data)
69

70
    return data, scale
71