Update app.py

app.py

@@ -184,24 +184,26 @@ def laplacian_blend(img1: np.ndarray, img2: np.ndarray, mask: np.ndarray, depth:
         
     return np.clip(imgs[-1], 0, 1)
 
-def get_image(img_path: str, mask: bool=False, scale: bool=True) -> np.array:
-    """
-    Gets image in appropriate format
-    """
-    if isinstance(img_path, np.ndarray):
-        img = img_path
+
+
+def get_image(img_input, mask=False, scale=True):
+    if isinstance(img_input, dict) and 'composite' in img_input:
+        img = img_input['composite']
+    elif isinstance(img_input, np.ndarray):
+        img = img_input
+    elif isinstance(img_input, str):
+        img = cv2.imread(img_input)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     else:
-        img = cv2.imread(img_path)
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB for file inputs
-    
+        raise ValueError("Unsupported image input type")
+
     if mask:
         if len(img.shape) == 3:
             img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-        _, binary_mask = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
-        return np.where(binary_mask == 255, 1, 0)
+        return np.where(img > 127, 1, 0)  # Threshold at 127 for the mask
     
-    if scale:
-        return img.astype('double') / 255.0
+    if scale and img.dtype != np.float64:
+        return img.astype('float64') / 255.0
     
     return img
 
@@ -253,11 +255,11 @@ with gr.Blocks(theme='bethecloud/storj_theme') as iface:
     )
     
     gr.Examples(
-        examples=[
-            ["img1.jpg", "img2.jpg", "mask1.jpg", "Poisson"],
-            ["img3.jpg", "img4.jpg", "mask2.jpg", "Mixed Gradient"],
-            ["img6.jpg", "img9.jpg", "mask3.jpg", "Laplacian"]
-        ],
+    examples=[
+        ["img1.jpg", "img2.jpg", {"composite": "mask1.jpg"}, "Poisson"],
+        ["img3.jpg", "img4.jpg", {"composite": "mask2.jpg"}, "Mixed Gradient"],
+        ["img6.jpg", "img9.jpg", {"composite": "mask3.jpg"}, "Laplacian"]
+    ],
         inputs=[bg_img, obj_img, mask_img, method],
         outputs=output_image,
         fn=blend_images,