jhj0517 commited on
Commit
87a101a
1 Parent(s): 002d880

Try except during inference

Browse files
Files changed (1) hide show
  1. modules/sam_inference.py +12 -5
modules/sam_inference.py CHANGED
@@ -74,7 +74,11 @@ class SamInference:
74
  model=self.model,
75
  **params
76
  )
77
- return self.mask_generator.generate(image)
 
 
 
 
78
 
79
  def predict_image(self,
80
  image: np.ndarray,
@@ -87,10 +91,13 @@ class SamInference:
87
  self.image_predictor = SAM2ImagePredictor(sam_model=self.model)
88
  self.image_predictor.set_image(image)
89
 
90
- masks, scores, logits = self.image_predictor.predict(
91
- box=box,
92
- multimask_output=params["multimask_output"],
93
- )
 
 
 
94
  return masks, scores, logits
95
 
96
  def divide_layer(self,
 
74
  model=self.model,
75
  **params
76
  )
77
+ try:
78
+ generated_masks = self.mask_generator.generate(image)
79
+ except Exception as e:
80
+ raise f"Error while auto generating masks: {e}"
81
+ return generated_masks
82
 
83
  def predict_image(self,
84
  image: np.ndarray,
 
91
  self.image_predictor = SAM2ImagePredictor(sam_model=self.model)
92
  self.image_predictor.set_image(image)
93
 
94
+ try:
95
+ masks, scores, logits = self.image_predictor.predict(
96
+ box=box,
97
+ multimask_output=params["multimask_output"],
98
+ )
99
+ except Exception as e:
100
+ raise f"Error while predicting image with prompt: {e}"
101
  return masks, scores, logits
102
 
103
  def divide_layer(self,