Update app.py

app.py

@@ -1,3 +1,9 @@
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import FileResponse
+from fastapi.responses import HTMLResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
+from typing import Tuple
+import cv2
 import fastai
 from fastai.vision import *
 from fastai.utils.mem import *
@@ -21,6 +27,8 @@ from PIL import Image
 from io import BytesIO
 import torchvision.transforms as T
 
+app = FastAPI()
+
 class FeatureLoss(nn.Module):
     def __init__(self, m_feat, layer_ids, layer_wgts):
         super().__init__()
@@ -94,10 +102,33 @@ def predict(img):
 
   return res, output_file
 
-gr.Interface(predict, 
-             inputs="image", 
-             outputs=[gr.Image(label="Sketch Image",show_share_button=False), gr.File(label="Result File")],
-             title = "<span style='color: crimson;'>Aiconvert.online</span>",
-             css="footer{display:none !important;}",
-             theme=gr.themes.Base(), 
-             description="").launch()
+@app.post("/predict/")
+async def predict(file: UploadFile = File(...)) -> Tuple[str, bytes]:
+    contents = await file.read()
+    img = cv2.imdecode(np.fromstring(contents, np.uint8), cv2.IMREAD_COLOR)
+    img = PIL.Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+    img = add_margin(img, 250, 250, 250, 250, (255, 255, 255))
+    img = np.array(img)
+
+    h, w = img.shape[:-1]
+    cv2.imwrite("test.jpg", img)
+    img_test = open_image("test.jpg")
+
+    p,img_hr,b = learn.predict(img_test)
+
+    res = (img_hr / img_hr.max()).numpy()
+    res = res[0]  # take only first channel as result
+    res = cv2.resize(res, (w,h))
+
+    output_file = get_filename()
+
+    cv2.imwrite(output_file, (res * 255).astype(np.uint8), [cv2.IMWRITE_JPEG_QUALITY, 50])
+
+    return output_file, res.tobytes()
+
+app.mount("/", StaticFiles(directory="static", html=True), name="static")
+
+@app.get("/")
+def index() -> FileResponse:
+    return FileResponse(path="/app/static/index.html", media_type="text/html")
+