Hector Lopez commited on
Commit
8a86d96
·
1 Parent(s): e5bb367
Files changed (1) hide show
  1. model.py +5 -2
model.py CHANGED
@@ -39,8 +39,11 @@ def get_checkpoint(checkpoint_path):
39
  return fixed_state_dict
40
 
41
  def predict(model, image, detection_threshold):
42
- img = PIL.Image.open(image)
43
- #img = PIL.Image.open(BytesIO(image))
 
 
 
44
  img = np.array(img)
45
  img = PIL.Image.fromarray(img)
46
 
 
39
  return fixed_state_dict
40
 
41
  def predict(model, image, detection_threshold):
42
+ if isinstance(image, str):
43
+ img = PIL.Image.open(image)
44
+ else:
45
+ img = PIL.Image.open(BytesIO(image))
46
+
47
  img = np.array(img)
48
  img = PIL.Image.fromarray(img)
49