kevinconka commited on
Commit
79e2714
·
1 Parent(s): fb2817b

better lazy loading via class

Browse files
Files changed (1) hide show
  1. app.py +17 -10
app.py CHANGED
@@ -50,20 +50,27 @@ h1 {
50
 
51
  rng = np.random.default_rng(0xDEADBEEF)
52
 
53
- model: AHOY = None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
54
 
55
  @spaces.GPU
56
  def inference(image):
57
  """Run inference on image and return annotated image."""
58
- global model
59
- if model is None:
60
- # lazy load model
61
- model = load_model("ahoy-RGB-b2")
62
- results = model(image)
63
- return results.draw(image, diameter=4)
64
-
65
-
66
- model = load_ahoy()
67
 
68
  # Flagging
69
  dataset_name = "SEA-AI/crowdsourced-sea-images"
 
50
 
51
  rng = np.random.default_rng(0xDEADBEEF)
52
 
53
+ class SeaVisionModel:
54
+ def __init__(self, model_name="ahoy-RGB-b2"):
55
+ self.model_name = model_name
56
+ self.model = None
57
+
58
+ def get_model(self):
59
+ if self.model is None:
60
+ self.model = load_model(self.model_name)
61
+ return self.model
62
+
63
+ def run_inference(self, image):
64
+ model = self.get_model()
65
+ results = model(image)
66
+ return results.draw(image, diameter=4)
67
+
68
+ vision_model = SeaVisionModel()
69
 
70
  @spaces.GPU
71
  def inference(image):
72
  """Run inference on image and return annotated image."""
73
+ return vision_model.run_inference(image)
 
 
 
 
 
 
 
 
74
 
75
  # Flagging
76
  dataset_name = "SEA-AI/crowdsourced-sea-images"