Enhance mask handling in image processing: resize and validate dimensions, update demo link, and correct copyright year in footer

src/core.py

@@ -400,13 +400,29 @@ def run(image, mask):
     return: BGR IMAGE
     """
     origin_height, origin_width = image.shape[1:]
+    
+    # Ensure mask has the same spatial dimensions as the image
+    if mask.shape[1:] != image.shape[1:]:
+        # Resize mask to match image dimensions
+        mask_resized = np.zeros((1, origin_height, origin_width), dtype=mask.dtype)
+        mask_resized[0, :, :] = cv2.resize(mask[0], (origin_width, origin_height), 
+                                          interpolation=cv2.INTER_NEAREST)
+        mask = mask_resized
+    
+    # Make sure both have dimensions divisible by 8
     image = pad_img_to_modulo(image, mod=8)
     mask = pad_img_to_modulo(mask, mod=8)
 
+    # Double check dimensions match after padding
+    if image.shape[1:] != mask.shape[1:]:
+        raise ValueError(f"Image and mask dimensions don't match after padding: {image.shape} vs {mask.shape}")
+    
     mask = (mask > 0) * 1
     image = torch.from_numpy(image).unsqueeze(0).to(device)
     mask = torch.from_numpy(mask).unsqueeze(0).to(device)
 
+    print(f"Input shapes - Image: {image.shape}, Mask: {mask.shape}")
+    
     start = time.time()
     with torch.no_grad():
         inpainted_image = model(image, mask)
@@ -432,21 +448,31 @@ def process_inpaint(image, mask):
     original_shape = image.shape
     interpolation = cv2.INTER_CUBIC
 
-    #size_limit: Union[int, str] = request.form.get("sizeLimit", "1080")
-    #if size_limit == "Original":
     size_limit = max(image.shape)
-    #else:
-    #    size_limit = int(size_limit)
 
     print(f"Origin image shape: {original_shape}")
     image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation)
     print(f"Resized image shape: {image.shape}")
+    
+    # Get alpha channel from mask
+    if mask.ndim == 3 and mask.shape[2] == 4:
+        mask = 255-mask[:,:,3]
+    elif mask.ndim == 3 and mask.shape[2] == 1:
+        mask = mask[:,:,0]
+    elif mask.ndim == 2:
+        mask = mask
+    else:
+        raise ValueError(f"Unexpected mask shape: {mask.shape}")
+    
+    # Ensure mask has the same dimensions as the image
+    if mask.shape[:2] != image.shape[:2]:
+        mask = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
+    
     image = norm_img(image)
-
-    mask = 255-mask[:,:,3]
-    mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation)
     mask = norm_img(mask)
-
+    
+    print(f"Normalized shapes - Image: {image.shape}, Mask: {mask.shape}")
+    
     res_np_img = run(image, mask)
 
     return cv2.cvtColor(res_np_img, cv2.COLOR_BGR2RGB)

templates/index.html

@@ -305,13 +305,13 @@
 
     <section class="demo-section">
         <h2>Try It Out</h2>
-        <p>Visit our interactive demo frontend:</p>
-        <a href="https://walidadebayo-magic-eraser.netlify.app" class="btn" target="_blank">Open Demo Frontend</a>
+        <p>Visit InnoAI interactive demo frontend:</p>
+        <a href="https://huggingface.co/spaces/innoai/Magic" class="btn" target="_blank">Open Demo Frontend</a>
         <p><small>Demo frontend connects to this API for processing images.</small></p>
     </section>
 
     <footer>
-        <p><small>© 2023 Magic Eraser API. MIT License.</small></p>
+        <p><small>© 2025 Magic Eraser API. MIT License.</small></p>
     </footer>
 </body>
 </html>