warm up model

Dockerfile

@@ -13,4 +13,7 @@ WORKDIR $HOME/app
 
 COPY --chown=user . $HOME/app
 
+# cache model
+RUN python3 warm_up.py
+
 CMD ["sh", "-c", "python3 main.py"]

dbimutils.py

@@ -65,5 +65,6 @@ def read_img_from_url(image):
             rawimage = Image.open(io.BytesIO(res.content))
             return rawimage
         else:
-            raise Exception("Invalid image type")
+            res = Image.open(image)
+            return rawimage
     return image

warm_up.py

@@ -0,0 +1,17 @@
+import img_label
+from img_nsfw import init_nsfw_pipe, check_nsfw
+
+if __name__ == "__main__":
+    img = 'test.jpeg'
+    pipe = init_nsfw_pipe()
+    nsfw_tags = check_nsfw(img, pipe)
+
+    score_general_threshold = 0.35
+    score_character_threshold = 0.85
+
+    img_tags = img_label.label_img(
+        image=img,
+        model="SwinV2",
+        l_score_general_threshold=score_general_threshold,
+        l_score_character_threshold=score_character_threshold,
+    )['general_res']