jini1114 commited on
Commit
c1ffcb6
1 Parent(s): 8cf0096

modified - TestDataset.py; include .exr in img_fmts

Browse files

modified - imutils.py; add opencv exr enable envirionment to consider .exr
; add imread for .exr image
; modified condition from 16-bit check to non 8-bit check

lib/common/imutils.py CHANGED
@@ -1,3 +1,5 @@
 
 
1
  import cv2
2
  import mediapipe as mp
3
  import torch
@@ -51,10 +53,13 @@ def get_affine_matrix_box(boxes, w2, h2):
51
 
52
  def load_img(img_file):
53
 
54
- img = cv2.imread(img_file, cv2.IMREAD_UNCHANGED)
 
 
 
55
 
56
- # considering 16-bit image
57
- if img.dtype == np.uint16:
58
  img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
59
 
60
  if len(img.shape) == 2:
 
1
+ import os
2
+ os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
3
  import cv2
4
  import mediapipe as mp
5
  import torch
 
53
 
54
  def load_img(img_file):
55
 
56
+ if img_file.endswith("exr"):
57
+ img = cv2.imread(img_file, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
58
+ else :
59
+ img = cv2.imread(img_file, cv2.IMREAD_UNCHANGED)
60
 
61
+ # considering non 8-bit image
62
+ if img.dtype != np.uint8 :
63
  img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
64
 
65
  if len(img.shape) == 2:
lib/dataset/TestDataset.py CHANGED
@@ -64,7 +64,7 @@ class TestDataset:
64
  self.device = device
65
 
66
  keep_lst = sorted(glob.glob(f"{self.image_dir}/*"))
67
- img_fmts = ["jpg", "png", "jpeg", "JPG", "bmp"]
68
 
69
  self.subject_list = sorted(
70
  [item for item in keep_lst if item.split(".")[-1] in img_fmts], reverse=False
 
64
  self.device = device
65
 
66
  keep_lst = sorted(glob.glob(f"{self.image_dir}/*"))
67
+ img_fmts = ["jpg", "png", "jpeg", "JPG", "bmp", "exr"]
68
 
69
  self.subject_list = sorted(
70
  [item for item in keep_lst if item.split(".")[-1] in img_fmts], reverse=False