feat: add support for different models

mbari-org · Sep 6, 2024 · 53743ff · 53743ff
1 parent 7902968
commit 53743ff
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 4 deletions.
diff --git a/.dockerignore b/.dockerignore
@@ -8,3 +8,5 @@ Dockerfile*
 *.md
 bin/
 src/submodules
+.github/
+.gitmodules
diff --git a/Dockerfile b/Dockerfile
@@ -1,7 +1,7 @@
 # ================================================================
 #  Docker image for fastapi-vss
 #  ================================================================
-FROM mbari/aidata:1.17.0
+FROM mbari/aidata:1.18.0
 
 LABEL vendor="MBARI"
 LABEL maintainer="dcline@mbari.org"

diff --git a/Dockerfile.cuda b/Dockerfile.cuda
@@ -1,7 +1,7 @@
 # ================================================================
 #  Docker image for fastapi-vss (CUDA enabled)
 #  ================================================================
-FROM mbari/aidata:1.17.0-cuda124
+FROM mbari/aidata:1.18.0-cuda124
 
 LABEL vendor="MBARI"
 LABEL maintainer="dcline@mbari.org"

diff --git a/src/app/config/__init__.py b/src/app/config/__init__.py
@@ -51,11 +51,12 @@ def init_config() -> dict:
             info(data)
             redis_host = data['redis']['host']
             redis_port = data['redis']['port']
+            model = data['vs']['model']
             password = os.getenv("REDIS_PASSWD")
 
             info(f"Connecting to redis at {redis_host}:{redis_port}")
             r = redis.Redis(host=redis_host, port=redis_port, password=password)
-            v = ViTWrapper(r, "cpu", False, BATCH_SIZE)
+            v = ViTWrapper(r, device="cpu", model_name=model, reset=False, batch_size=BATCH_SIZE)
 
             project = data['tator']['project']
 

diff --git a/src/app/project/config/config.yml b/src/app/project/config/config.yml
@@ -17,7 +17,8 @@ sdcat:
 
 vss:
   project: "cfe"
-  threshold: 0.4
+  threshold: 0.9
+  model: "google/vit-base-patch16-224"
 
 tator:
   project: "902111-CFE"