svjack commited on
Commit
e435229
·
1 Parent(s): bfc3012

Update pred_color.py

Browse files
Files changed (1) hide show
  1. pred_color.py +2 -0
pred_color.py CHANGED
@@ -94,6 +94,8 @@ import os
94
  def single_pred_features(image):
95
  if type(image) == type("") and os.path.exists(image):
96
  image = Image.open(image).convert("RGB")
 
 
97
  else:
98
  image = image.convert("RGB")
99
  cv2_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
94
  def single_pred_features(image):
95
  if type(image) == type("") and os.path.exists(image):
96
  image = Image.open(image).convert("RGB")
97
+ elif hasattr(image, "shape"):
98
+ image = Image.fromarray(image).convert("RGB")
99
  else:
100
  image = image.convert("RGB")
101
  cv2_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)