Thanusha/Image_processing

1from PIL import Image2import numpy as np3import cv24from rembg import remove5 6 7# Threshold Variables8BLUR_THRESHOLD = 100  # Variance of Laplacian threshold. Typical values: min=10, avg=100, max=2009EDGE_COUNT_THRESHOLD = 2000  # Edge count threshold. Typical values: min=500, avg=1000, max=200010WHITE_BG_THRESHOLD = 240  # RGB threshold to consider a pixel as white. Typical values: min=200, avg=240, max=25511WHITE_BG_PERCENTAGE = 0.9  # Percentage of white pixels in the corners to consider background as white. Typical values: min=0.5, avg=0.9, max=1.012 13def crop_to_content_with_alpha(image):14    """Crop image to content based on tranperancy for alpha channel"""15    if image.mode != 'RGBA':16        image = image.convert('RGBA')17    alpha_channel = image.split()[3]18    if alpha_channel.getextrema()[0]<255:19        bbox = alpha_channel.getbbox()20        if bbox:21            cropped_image = image.crop(bbox)22            return cropped_image23    return image24 25def crop_to_content_without_alpha(image_path):26    """Crop image to content based on contours for images without alpha channel"""27    image = Image.open(image_path)28    open_cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGRA)29    gray = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2GRAY)30    blurred = cv2.GaussianBlur(gray,(5,5), 0)31    edges = cv2.Canny(blurred, 50, 150)32    dilated = cv2.dilate(edges, None, iterations=2)33    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)34    if contours:35        x, y, w, h = cv2.boundingRect(contours[0])36        cropped_image = open_cv_image[y:y+h, x:x+w]37        cropped_image_rgb = cv2.cvtColor(cropped_image, cv2.COLOR_BAYER_BG2BGR)38        pil_cropped_image = Image.fromarray(cropped_image_rgb)39        return pil_cropped_image40    return Image.open(image_path)41 42def is_image_vertical(image):43    """Check if the image us vertical (portrait orientation)"""44    width, height = image.size45    return height>width46 47def rotate_to_vertical(image):48    "Rotate image is vertical if it is not portrait oriented"49    if not is_image_vertical(image):50        image = image.rotate(90, expand=True)51    return image52 53def remove_background(img):54    """Remove the background from the image"""55    try:56        output = remove(img)57        return output58    except Exception as e:59        raise RuntimeError(f"Error removing background: {e}")60    61def resize_image(cropped_image, canvas_width, canvas_height, fixed_margin):62    """ Resize the image to fit within the specified canvas dimensions"""63    aspect_ratio = cropped_image.width / cropped_image.height64    max_width = canvas_width -2 * fixed_margin65    max_height = canvas_height -2 * fixed_margin66    scale_factor = min(max_width/ cropped_image.width, max_height / cropped_image.height)67    new_width = int(cropped_image.width * scale_factor)68    new_height = int(cropped_image.height * scale_factor)69    resized_image = cropped_image.resize((new_width, new_height), Image.LANCZOS)70    paste_x = (canvas_width - new_width) //271    paste_y = (canvas_height - new_height)//272    final_image = Image.new("RGBA", (canvas_height, canvas_width), (0,0,0,0))73    # Fixed variable names and mask handling74    mask = resized_image.split()[3] if "A" in resized_image.getbands() else None75    final_image.paste(resized_image, (paste_x, paste_y), mask=mask)76    return final_image77    78def process_image(image, canvas_width, canvas_height, fixed_margin):79    """Process a single image: removes background , crop , rotate, and resize"""80    output = remove_background(image)81    if output:82        cropped_image = crop_to_content_with_alpha(output)83        cropped_image = rotate_to_vertical(cropped_image)84        resized_image = resize_image(cropped_image, canvas_width, canvas_height, fixed_margin)85        return resized_image86    return None