import os from PIL import Image import io from rembg import remove from transformers import pipeline import numpy as np def remove_background_rembg(input_path): print(f"Removing background using rembg for image: {input_path}") with open(input_path, 'rb') as i: input_image = i.read() output_image = remove(input_image) img = Image.open(io.BytesIO(output_image)).convert("RGBA") return img def remove_background_bria(input_path): print(f"Removing background using bria for image: {input_path}") pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True) pillow_image = pipe(input_path) return pillow_image ###### PlACE TO PUT ANOTHER MODEL ####### def get_bounding_box_with_threshold(image, threshold): # Convert image to numpy array img_array = np.array(image) # Get alpha channel alpha = img_array[:,:,3] # Find rows and columns where alpha > threshold rows = np.any(alpha > threshold, axis=1) cols = np.any(alpha > threshold, axis=0) # Find the bounding box top, bottom = np.where(rows)[0][[0, -1]] left, right = np.where(cols)[0][[0, -1]] if left < right and top < bottom: return (left, top, right, bottom) else: return None def position_logic(image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left, use_threshold=True): image = Image.open(image_path) image = image.convert("RGBA") # Get the bounding box of the non-blank area with threshold if use_threshold: bbox = get_bounding_box_with_threshold(image, threshold=10) else: bbox = image.getbbox() log = [] if bbox: # Check 1 pixel around the image for non-transparent pixels width, height = image.size cropped_sides = [] # Define tolerance for transparency tolerance = 30 # Adjust this value as needed # Check top edge if any(image.getpixel((x, 0))[3] > tolerance for x in range(width)): cropped_sides.append("top") # Check bottom edge if any(image.getpixel((x, height-1))[3] > tolerance for x in range(width)): cropped_sides.append("bottom") # Check left edge if any(image.getpixel((0, y))[3] > tolerance for y in range(height)): cropped_sides.append("left") # Check right edge if any(image.getpixel((width-1, y))[3] > tolerance for y in range(height)): cropped_sides.append("right") if cropped_sides: info_message = f"Info for {os.path.basename(image_path)}: The following sides of the image may contain cropped objects: {', '.join(cropped_sides)}" print(info_message) log.append({"info": info_message}) else: info_message = f"Info for {os.path.basename(image_path)}: The image is not cropped." print(info_message) log.append({"info": info_message}) # Crop the image to the bounding box image = image.crop(bbox) log.append({"action": "crop", "bbox": [str(bbox[0]), str(bbox[1]), str(bbox[2]), str(bbox[3])]}) # Calculate the new size to expand the image target_width, target_height = canvas_size aspect_ratio = image.width / image.height if len(cropped_sides) == 4: # If the image is cropped on all sides, center crop it to fit the canvas if aspect_ratio > 1: # Landscape new_height = target_height new_width = int(new_height * aspect_ratio) left = (new_width - target_width) // 2 image = image.resize((new_width, new_height), Image.LANCZOS) image = image.crop((left, 0, left + target_width, target_height)) else: # Portrait or square new_width = target_width new_height = int(new_width / aspect_ratio) top = (new_height - target_height) // 2 image = image.resize((new_width, new_height), Image.LANCZOS) image = image.crop((0, top, target_width, top + target_height)) log.append({"action": "center_crop_resize", "new_size": f"{target_width}x{target_height}"}) x, y = 0, 0 elif not cropped_sides: # If the image is not cropped, expand it from center until it touches the padding new_height = target_height - padding_top - padding_bottom new_width = int(new_height * aspect_ratio) if new_width > target_width - padding_left - padding_right: # If width exceeds available space, adjust based on width new_width = target_width - padding_left - padding_right new_height = int(new_width / aspect_ratio) # Resize the image image = image.resize((new_width, new_height), Image.LANCZOS) log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)}) x = (target_width - new_width) // 2 y = target_height - new_height - padding_bottom else: # New logic for handling cropped top and left, or top and right if set(cropped_sides) == {"top", "left"} or set(cropped_sides) == {"top", "right"}: new_height = target_height - padding_bottom new_width = int(new_height * aspect_ratio) # If new width exceeds canvas width, adjust based on width if new_width > target_width: new_width = target_width new_height = int(new_width / aspect_ratio) # Resize the image image = image.resize((new_width, new_height), Image.LANCZOS) log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)}) # Set position if "left" in cropped_sides: x = 0 else: # right in cropped_sides x = target_width - new_width y = 0 # If the resized image is taller than the canvas minus padding, crop from the bottom if new_height > target_height - padding_bottom: crop_bottom = new_height - (target_height - padding_bottom) image = image.crop((0, 0, new_width, new_height - crop_bottom)) new_height = target_height - padding_bottom log.append({"action": "crop_vertical", "bottom_pixels_removed": str(crop_bottom)}) log.append({"action": "position", "x": str(x), "y": str(y)}) elif set(cropped_sides) == {"bottom", "left"} or set(cropped_sides) == {"bottom", "right"}: # Handle bottom & left or bottom & right cropped images new_height = target_height - padding_top new_width = int(new_height * aspect_ratio) # If new width exceeds canvas width, adjust based on width if new_width > target_width - padding_left - padding_right: new_width = target_width - padding_left - padding_right new_height = int(new_width / aspect_ratio) # Resize the image without cropping or stretching image = image.resize((new_width, new_height), Image.LANCZOS) log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)}) # Set position if "left" in cropped_sides: x = 0 else: # right in cropped_sides x = target_width - new_width y = target_height - new_height log.append({"action": "position", "x": str(x), "y": str(y)}) elif set(cropped_sides) == {"bottom", "left", "right"}: # Expand the image from the center new_width = target_width new_height = int(new_width / aspect_ratio) if new_height < target_height: new_height = target_height new_width = int(new_height * aspect_ratio) image = image.resize((new_width, new_height), Image.LANCZOS) # Crop to fit the canvas left = (new_width - target_width) // 2 top = 0 image = image.crop((left, top, left + target_width, top + target_height)) log.append({"action": "expand_and_crop", "new_size": f"{target_width}x{target_height}"}) x, y = 0, 0 elif cropped_sides == ["top"]: # New logic for handling only top-cropped images if image.width > image.height: new_width = target_width new_height = int(target_width / aspect_ratio) else: new_height = target_height - padding_bottom new_width = int(new_height * aspect_ratio) # Resize the image image = image.resize((new_width, new_height), Image.LANCZOS) log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)}) x = (target_width - new_width) // 2 y = 0 # Align to top # Apply padding only to non-cropped sides x = max(padding_left, min(x, target_width - new_width - padding_right)) elif cropped_sides in [["right"], ["left"]]: # New logic for handling only right-cropped or left-cropped images if image.width > image.height: new_width = target_width - max(padding_left, padding_right) new_height = int(new_width / aspect_ratio) else: new_height = target_height - padding_top - padding_bottom new_width = int(new_height * aspect_ratio) # Resize the image image = image.resize((new_width, new_height), Image.LANCZOS) log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)}) if cropped_sides == ["right"]: x = target_width - new_width # Align to right else: # cropped_sides == ["left"] x = 0 # Align to left y = target_height - new_height - padding_bottom # Respect bottom padding # Ensure top padding is respected if y < padding_top: y = padding_top log.append({"action": "position", "x": str(x), "y": str(y)}) elif set(cropped_sides) == {"left", "right"}: # Logic for handling images cropped on both left and right sides new_width = target_width # Expand to full width of canvas # Calculate the aspect ratio of the original image aspect_ratio = image.width / image.height # Calculate the new height while maintaining aspect ratio new_height = int(new_width / aspect_ratio) # Resize the image image = image.resize((new_width, new_height), Image.LANCZOS) log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)}) # Set horizontal position (always 0 as it spans full width) x = 0 # Calculate vertical position to respect bottom padding y = target_height - new_height - padding_bottom # If the resized image is taller than the canvas, crop from the top only if new_height > target_height - padding_bottom: crop_top = new_height - (target_height - padding_bottom) image = image.crop((0, crop_top, new_width, new_height)) new_height = target_height - padding_bottom y = 0 log.append({"action": "crop_vertical", "top_pixels_removed": str(crop_top)}) else: # Align the image to the bottom with padding y = target_height - new_height - padding_bottom log.append({"action": "position", "x": str(x), "y": str(y)}) elif cropped_sides == ["bottom"]: # Logic for handling images cropped on the bottom side # Calculate the aspect ratio of the original image aspect_ratio = image.width / image.height if aspect_ratio < 1: # Portrait orientation new_height = target_height - padding_top # Full height with top padding new_width = int(new_height * aspect_ratio) # If the new width exceeds the canvas width, adjust it if new_width > target_width: new_width = target_width new_height = int(new_width / aspect_ratio) else: # Landscape orientation new_width = target_width - padding_left - padding_right new_height = int(new_width / aspect_ratio) # If the new height exceeds the canvas height, adjust it if new_height > target_height: new_height = target_height new_width = int(new_height * aspect_ratio) # Resize the image image = image.resize((new_width, new_height), Image.LANCZOS) log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)}) # Set horizontal position (centered) x = (target_width - new_width) // 2 # Set vertical position (touching bottom edge for all cases) y = target_height - new_height log.append({"action": "position", "x": str(x), "y": str(y)}) else: # Use the original resizing logic for other partially cropped images if image.width > image.height: new_width = target_width new_height = int(target_width / aspect_ratio) else: new_height = target_height new_width = int(target_height * aspect_ratio) # Resize the image image = image.resize((new_width, new_height), Image.LANCZOS) log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)}) # Center horizontally for all images x = (target_width - new_width) // 2 y = target_height - new_height - padding_bottom # Adjust positions for cropped sides if "top" in cropped_sides: y = 0 elif "bottom" in cropped_sides: y = target_height - new_height if "left" in cropped_sides: x = 0 elif "right" in cropped_sides: x = target_width - new_width # Apply padding only to non-cropped sides, but keep horizontal centering if "left" not in cropped_sides and "right" not in cropped_sides: x = (target_width - new_width) // 2 # Always center horizontally if "top" not in cropped_sides and "bottom" not in cropped_sides: y = max(padding_top, min(y, target_height - new_height - padding_bottom)) return log, image, x, y