Spaces:
Running
Running
Commit
·
79e2714
1
Parent(s):
fb2817b
better lazy loading via class
Browse files
app.py
CHANGED
@@ -50,20 +50,27 @@ h1 {
|
|
50 |
|
51 |
rng = np.random.default_rng(0xDEADBEEF)
|
52 |
|
53 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
54 |
|
55 |
@spaces.GPU
|
56 |
def inference(image):
|
57 |
"""Run inference on image and return annotated image."""
|
58 |
-
|
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"
|