Update app.py

app.py

@@ -7,7 +7,7 @@ import cv2
 import numpy as np
 
 from camera_input_live import camera_input_live
-
+from inference import inpainting
 st.set_page_config(layout="wide")
 
 
@@ -40,7 +40,9 @@ def get_mask(image_mask: np.ndarray) -> np.ndarray:
     if mask.sum() > 0:
         mask = mask * 1
     # 3 channels
-    mask = np.stack([mask, mask, mask], axis=2)
+    mask = np.stack([mask, mask, mask], axis=2)*255
+    mask = mask.astype(np.uint8)
+    mask = Image.fromarray(mask).convert("RGB")
     return mask
 
 
@@ -69,7 +71,7 @@ def make_input_fields():
 def decode_image(image):
     cv2_img = cv2.imdecode(np.frombuffer(image, np.uint8), cv2.IMREAD_COLOR)
     cv2_img = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)
-    image = Image.fromarray(cv2_img).convert("RGB")
+    image = Image.fromarray(cv2_img).convert("RGB").resize((512, 512))
     return image
 
 if __name__ == "__main__":
@@ -91,12 +93,17 @@ if __name__ == "__main__":
         image = decode_image(webcam.getvalue())
 
         canvas = make_canvas(image)
-
+    
+    if st.button("Inpaint"):
         mask = get_mask(np.array(canvas.image_data))
+        result = inpainting(image, mask_image, prompt, negative_prompt)
+    else:
+        result = None
 
     with colB:
         st.write("## Generated image")
         st.write("The generated image will appear here.")
         if webcam:
             st.image(webcam)
-        # st.image(mask*255)
+        if result:
+            st.image(result)