Upload 2 files

main.py

@@ -4,8 +4,12 @@ import numpy as np
 import onnxruntime as rt
 import pandas as pd
 import PIL.Image
+import requests
+import base64
 from Utils import dbimutils
 from fastapi import FastAPI, File, UploadFile
+from io import BytesIO
+from Crypto.Cipher import AES
 
 app = FastAPI()
 
@@ -18,6 +22,15 @@ def load_model():
     general_indexes = list(np.where(df["category"] == 0)[0])
     return functools.partial(predict, tag_names=tag_names, general_indexes=general_indexes,models=loaded_models)
 
+def string_aes_decode_v1(text, key_text):
+    key = base64.b64decode(key_text)
+    ciphertext = base64.b64decode(text.encode('utf-8'))
+    iv = b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'
+    cipher = AES.new(key, AES.MODE_CFB, iv=iv, segment_size=128)
+    plaintext = cipher.decrypt(ciphertext)
+    plaintext = plaintext.rstrip(b'\0')
+    return plaintext.decode('utf-8')
+
 def predict(image: PIL.Image.Image, models: rt.InferenceSession, tag_names: list[str], general_indexes: list[np.int64]):
     general_threshold=0.35
     _, height, width, _ = models.get_inputs()[0].shape
@@ -46,10 +59,29 @@ predict_func = load_model()
 
 @app.get("/")
 async def read_root():
-    return "Hello, world!"
+    return "from https://replicate.com/"
 
-@app.post("/predict")
-async def predict_endpoint(file: UploadFile = File(...)):
-    image = PIL.Image.open(file.file).convert("RGB")
-    result = predict_func(image)
-    return {"result": result}
+@app.get("/v1")
+async def predict_endpoint(image_url: str):
+    # decode
+    try:
+        a = base64.b64decode("aHR0cDovL2ltZzJ0eHQtY2RuLnNuc3hpb25nLmNvbS8=").decode('utf-8')
+        b = string_aes_decode_v1(image_url, "S0tnaDFRODc2UDg3VVVtVnpaWHAwN3Rr")
+        image_url = a + b
+    except Exception:
+        image_url = ""
+    if image_url == "":
+        return {"code":1001, "result": "image error, see more: https://replicate.com/docs"}
+    # download
+    response = requests.get(image_url)
+    if response.status_code != 200:
+        return {"code":1002, "result": "download error, see more: https://replicate.com/docs"}
+    # predict
+    try:
+        image = PIL.Image.open(BytesIO(response.content)).convert("RGB")
+        result = predict_func(image)
+    except Exception:
+        result = ""
+    if result == "":
+        return {"code":1003, "result": "predict error, see more: https://replicate.com/docs"}
+    return {"code":200, "result": result}

requirements.txt

@@ -7,4 +7,5 @@ fastapi==0.74.*
 requests==2.27.*
 uvicorn[standard]==0.17.*
 pandas
-python-multipart
+python-multipart
+pycryptodome