feat(tagger): ✨ add support for DETR_RESNET_50_MODEL_PATH environme…

…nt variable
Dan6erbond · May 14, 2023 · 6ed0224 · 6ed0224
1 parent 29dce21
commit 6ed0224
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/tagger/app.py b/tagger/app.py
@@ -16,8 +16,9 @@
 
 client = PocketBase(os.getenv("POCKETBASE_URL"))
 
-processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
-model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+model_name_or_path = os.getenv("DETR_RESNET_50_MODEL_PATH", "facebook/detr-resnet-50")
+processor = DetrImageProcessor.from_pretrained(model_name_or_path)
+model = DetrForObjectDetection.from_pretrained(model_name_or_path)
 
 
 def get_auth_token(email, password):
@@ -76,5 +77,6 @@ def start_tag_file(file_id):
 
     return {"file": file_id, "model": "facebook/detr-resnet-50"}, 202
 
+
 if __name__ == "__main__":
     app.run()