Dockerfile order, basic test + fixture

gadicc · gadicc · commit 33cd2566b062 · 2022-11-03T06:59:38.000Z
diff --git a/Dockerfile b/Dockerfile
@@ -8,17 +8,18 @@ RUN apt-key adv --keyserver keyserver.ubuntu.com --recv-keys A4B469963BF863CC
 ENV DEBIAN_FRONTEND=noninteractive
 
 # Install git and OpenCV
-RUN apt-get update && apt-get install -yqq git libopencv-dev python3-opencv
+RUN apt-get update && apt-get install -yqq git
+
+WORKDIR /
+
+COPY root-cache/. /root/.cache
 
 # Compile PyPatchMatch
+RUN apt-get install -yqq libopencv-dev python3-opencv
 WORKDIR /PyPatchMatch
 ADD PyPatchMatch .
 RUN make
 
-WORKDIR /
-
-COPY root-cache/. /root/.cache
-
 # Install python packages
 RUN pip3 install --upgrade pip
 ADD requirements.txt requirements.txt
diff --git a/app.py b/app.py
@@ -20,6 +20,7 @@
 import numpy as np
 import skimage
 import skimage.measure
+import patch_match
 
 MODEL_ID = os.environ.get("MODEL_ID")
 PIPELINE = os.environ.get("PIPELINE")
@@ -217,7 +218,7 @@ def inference(all_inputs: dict) -> dict:
         sel_buffer = np.array(model_inputs.get("init_image"))
         img = sel_buffer[:, :, 0:3]
         mask = sel_buffer[:, :, -1]
-        img = patch_match.inpaint(img, mask=255-mask, patch_size=3)
+        img = patch_match.inpaint(img, mask=255 - mask, patch_size=3)
         model_inputs["init_image"] = PIL.Image.fromarray(img)
         mask = 255 - mask
         mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
diff --git a/test2.py b/test2.py
@@ -0,0 +1,70 @@
+# This file is used to verify your http server acts as expected
+# Run it with `python3 test.py``
+
+import requests
+import base64
+import os
+import json
+from io import BytesIO
+from PIL import Image
+from pathlib import Path
+
+TESTS = "tests"
+FIXTURES = TESTS + os.sep + "fixtures"
+OUTPUT = TESTS + os.sep + "output"
+Path(OUTPUT).mkdir(parents=True, exist_ok=True)
+
+
+def b64encode_file(filename: str):
+    with open(os.path.join(FIXTURES, filename), "rb") as file:
+        return base64.b64encode(file.read())
+
+
+def output_path(filename: str):
+    return os.path.join(OUTPUT, filename)
+
+
+def decode_and_save(image_byte_string: str, name: str):
+    image_encoded = image_byte_string.encode("utf-8")
+    image_bytes = BytesIO(base64.b64decode(image_encoded))
+    image = Image.open(image_bytes)
+    fp = output_path(name + ".png")
+    image.save(fp)
+    print("Saved " + fp)
+
+
+def test(name, inputs):
+    print("Running test: " + name)
+    response = requests.post("http://localhost:8000/", json=inputs)
+    result = response.json()
+
+    if result.get("images_base64", None) == "None":
+        print(json.dumps(result, indent=4))
+        print()
+        return
+
+    images_base64 = result.get("images_base64", None)
+    if images_base64:
+        for idx, image_byte_string in enumerate(images_base64):
+            decode_and_save(image_byte_string, f"{name}_{idx}")
+    else:
+        decode_and_save(result["image_base64"], name)
+
+    print()
+
+
+test(
+    "outpaint",
+    {
+        "modelInputs": {
+            "prompt": "girl with a pearl earing standing in a big room",
+            "init_image": b64encode_file("girl_with_pearl_earing_outpainting_in.png"),
+        },
+        "callInputs": {
+            "MODEL_ID": "CompVis/stable-diffusion-v1-4",
+            "PIPELINE": "StableDiffusionInpaintPipelineLegacy",
+            "SCHEDULER": "DDIM",  # Note, as of diffusers 0.3.0, no LMS yet
+            "FILL_MODE": "patchmatch",
+        },
+    },
+)
diff --git a/tests/fixtures/girl_with_pearl_earing_outpainting_in.png b/tests/fixtures/girl_with_pearl_earing_outpainting_in.png
