image_tools.py

# -*- coding: utf-8 -*-
"""image_tools.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1BoqBIyOSwHQ5j1mjZ6ScvQBgCCcGulNf
"""

# Image Tools Package
# Below is a complete Python package module implementing:
# 1. Background removal
# 2. Add background color/template
# 3. Enhance image quality
# 4. Resize and compress image
# 5. Resize by custom height*width
#
# This file can be saved as `image_tools.py` and packaged via setup.py or pyproject.toml
# Required libraries:
# - pillow
# - rembg (for background removal)
# - io, os
# - numpy (optional for quality enhancement)
#
# Installation example:
#   pip install pillow rembg numpy
#
# Usage example:
#   from image_tools import ImageProcessor
#   p = ImageProcessor()
#   p.remove_bg("input.jpg", "out.png")
#
# ---------------------------------------------------------------

#!pip install rembg pillow onnxruntime

import os
from PIL import Image, ImageEnhance
from rembg import remove
import io

class ImageProcessor:
    """
    A utility class providing:
    1. Background removal
    2. Background replacement
    3. Image enhancement
    4. Compression (reduce KB size)
    5. Resize with custom height/width
    """

    def __init__(self):
        pass

    # -----------------------------------------------------------
    # 1. REMOVE BACKGROUND
    # -----------------------------------------------------------
    def remove_bg(self, input_path: str, output_path: str):
        """Remove background using rembg library.
        Handles errors and supports PNG transparent output."""
        try:
            with open(input_path, "rb") as i:
                img_data = i.read()
                result = remove(img_data)

            with open(output_path, "wb") as o:
                o.write(result)

        except Exception as e:
            raise RuntimeError(f"Failed to remove background: {e}")

    # -----------------------------------------------------------
    # 2. ADD BACKGROUND COLOR OR TEMPLATE
    # -----------------------------------------------------------
    def add_background(self, input_path: str, output_path: str, bg_color=(255, 255, 255), template_path: str = None):
        """Add background color or apply a template.
        If template_path is provided, it overlays the template."""
        try:
            img = Image.open(input_path).convert("RGBA")

            if template_path:
                bg = Image.open(template_path).convert("RGBA").resize(img.size)
            else:
                bg = Image.new("RGBA", img.size, bg_color + (255,))

            combined = Image.alpha_composite(bg, img)

            combined.convert("RGB").save(output_path, quality=95)
        except Exception as e:
            raise RuntimeError(f"Failed to apply background: {e}")

    # -----------------------------------------------------------
    # 3. ENHANCE IMAGE QUALITY
    # -----------------------------------------------------------
    def enhance_quality(self, input_path: str, output_path: str, sharpness=1.5, contrast=1.3, brightness=1.1):
        """Enhances image using Pillow built-in enhancement modules."""
        try:
            img = Image.open(input_path).convert("RGB")

            # Sharpness
            img = ImageEnhance.Sharpness(img).enhance(sharpness)
            # Contrast
            img = ImageEnhance.Contrast(img).enhance(contrast)
            # Brightness
            img = ImageEnhance.Brightness(img).enhance(brightness)

            img.save(output_path, quality=95)
        except Exception as e:
            raise RuntimeError(f"Failed to enhance quality: {e}")

    # -----------------------------------------------------------
    # 4. RESIZE & COMPRESS TO KB RANGE
    # -----------------------------------------------------------
    def compress_to_size(self, input_path: str, output_path: str, target_kb: int):
        """Compress an image to a target KB size while preserving quality.
        Uses iterative reduction until size < target_kb."""
        try:
            img = Image.open(input_path)
            quality = 95

            # Save iteratively reducing quality
            while True:
                buffer = io.BytesIO()
                img.save(buffer, format="JPEG", quality=quality)
                size_kb = buffer.tell() / 1024

                if size_kb <= target_kb or quality <= 5:
                    with open(output_path, "wb") as f:
                        f.write(buffer.getvalue())
                    break

                quality -= 5  # reduce quality step-by-step

        except Exception as e:
            raise RuntimeError(f"Failed to compress image: {e}")

    # -----------------------------------------------------------
    # 5. CUSTOM RESIZE (HEIGHT * WIDTH)
    # -----------------------------------------------------------
    def resize_dimensions(self, input_path: str, output_path: str, width: int, height: int):
        """Resize image to custom dimensions with high-quality LANCZOS filter."""
        try:
            img = Image.open(input_path)
            resized = img.resize((width, height), Image.LANCZOS)
            resized.save(output_path, quality=95)
        except Exception as e:
            raise RuntimeError(f"Failed to resize image: {e}")



processor = ImageProcessor()
input_image_path = "/content/output.png"
output_image_path_with_bg = "output_with_bg.png"

bg_choice = input("Do you want to use a template image (T) or a solid background color (C)? (T/C): ").strip().upper()

if bg_choice == 'T':
    template_path = input("Please enter the path to your template image: ").strip()
    if not template_path:
        print("No template path provided. Using default white background.")
        processor.add_background(input_image_path, output_image_path_with_bg)
    else:
        processor.add_background(input_image_path, output_image_path_with_bg, template_path=template_path)
elif bg_choice == 'C':
    color_input = input("Enter background color as R,G,B (e.g., 255,255,255 for white) or a common color name (e.g., white, black, red, green, blue): ").strip().lower()
    bg_color = (255, 255, 255) # Default to white
    if color_input == 'white':
        bg_color = (255, 255, 255)
    elif color_input == 'black':
        bg_color = (0, 0, 0)
    elif color_input == 'red':
        bg_color = (255, 0, 0)
    elif color_input == 'green':
        bg_color = (0, 255, 0)
    elif color_input == 'blue':
        bg_color = (0, 0, 255)
    else:
        try:
            r, g, b = map(int, color_input.split(','))
            bg_color = (r, g, b)
        except ValueError:
            print("Invalid color format. Using default white background.")

    processor.add_background(input_image_path, output_image_path_with_bg, bg_color=bg_color)
else:
    print("Invalid choice. Using default white background.")
    processor.add_background(input_image_path, output_image_path_with_bg)

print(f"Background added to {input_image_path} and saved to {output_image_path_with_bg}")

processor = ImageProcessor()
input_image_path = "/content/input.jpeg"
output_image_path = "output.png"

processor.remove_bg(input_image_path, output_image_path)
print(f"Background removed from {input_image_path} and saved to {output_image_path}")