Update crop function

app.py

@@ -70,31 +70,42 @@ def build_model(hypar,device):
     return net
 
 
-def crop_to_signature(mask, padding=32):
+def crop_signature(original_image_path, mask, padding=32):
     """
-    Crop the signature area based on the mask and add padding.
+    Crop the signature from the original image using the provided mask.
+    
+    :param original_image_path: The file path of the original image.
     :param mask: The binary mask of the signature.
-    :param padding: Padding around the cropped signature.
-    :return: Cropped mask of the signature with padding.
+    :param padding: Padding to add around the bounding box of the signature.
+    :return: Cropped image containing the signature.
     """
-    contours, _ = cv2.findContours(
-        mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    # Convert the mask to a binary image
+    _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
+
+    # Find contours from the binary mask
+    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    
+    # Open the original image
+    original_image = Image.open(original_image_path).convert("RGB")
+    
+    # If contours are found, proceed to crop
     if contours:
         # Assume the largest contour is the signature
         x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea))
 
         # Add padding to the bounding box
-        x = max(x - padding, 0)
-        y = max(y - padding, 0)
-        w = min(w + 2 * padding, mask.shape[1] - x)
-        h = min(h + 2 * padding, mask.shape[0] - y)
-
-        # Crop the mask
-        cropped_mask = mask[y:y+h, x:x+w]
-        return cropped_mask
-    else:
-        # Return the original mask if no contours are found
-        return mask
+        x_padded = max(x - padding, 0)
+        y_padded = max(y - padding, 0)
+        w_padded = min(w + 2 * padding, original_image.width - x_padded)
+        h_padded = min(h + 2 * padding, original_image.height - y_padded)
+
+        # Crop the original image and mask using the bounding box with padding
+        cropped_image = original_image.crop((x_padded, y_padded, x_padded + w_padded, y_padded + h_padded))
+        
+        return cropped_image
+    
+    # If no contours are found, return the original image
+    return original_image
 
 
 def smooth_and_denoise(mask):
@@ -107,11 +118,11 @@ def smooth_and_denoise(mask):
     smoothed_mask = cv2.GaussianBlur(mask, (5, 5), 0)
 
     # Estimate noise standard deviation from the image
-    sigma_est = np.mean(estimate_sigma(smoothed_mask, channel_axis=None))
+    sigma_est = np.mean(estimate_sigma(smoothed_mask, channel_axis=-1))
 
     # Apply Non-Local Means Denoising
     denoised_mask = denoise_nl_means(smoothed_mask, h=1.15 * sigma_est, fast_mode=True,
-                                     patch_size=5, patch_distance=3, channel_axis=None)
+                                     patch_size=5, patch_distance=3, channel_axis=-1)
     return denoised_mask
 
     
@@ -166,28 +177,42 @@ hypar["model"] = ISNetDIS()
  # Build Model
 net = build_model(hypar, device)
 
+def inference(image):
+    image_path = image
+  
+    image_tensor, orig_size = load_image(image_path, hypar) 
+    mask = predict(net, image_tensor, orig_size, hypar, device)
+
+    cropped_mask = crop_to_signature(mask)
+    processed_mask = smooth_and_denoise(cropped_mask)
+
+    # Convert to PIL image for output
+    pil_mask = Image.fromarray((processed_mask * 255).astype(np.uint8)).convert('L')
+    return pil_mask
+
 
 def inference(image):
     image_path = image
     
     image_tensor, orig_size = load_image(image_path, hypar) 
     original_mask = predict(net, image_tensor, orig_size, hypar, device)
-    
+     
     # Process the original mask with smoothing and denoising
     processed_mask = smooth_and_denoise(original_mask)
 
     # Convert processed mask to PIL image
-    pil_processed_mask = Image.fromarray(processed_mask).convert('L')
+    pil_processed_mask = Image.fromarray((processed_mask * 255).astype(np.uint8)).convert('L')
+    
     im_rgb = Image.open(image).convert("RGB")
-
     im_dark = Image.new('RGB', im_rgb.size, (0, 0, 0))
+    cropped_signature_image = crop_signature(image_path, mask)
     
     # Apply processed mask to images
     im_rgba = im_rgb.copy()
-    im_rgba.putalpha(pil_processed_mask)
+    im_rgba.putalpha(original_mask)
     im_dark.putalpha(pil_processed_mask)
     
-    return [im_rgba, pil_processed_mask, im_dark]
+    return [cropped_signature_image, processed_mask, im_dark]
 
 
 title = "Mysign.id - Signature Background removal based on DIS"