| import numpy as np
|
| from PIL import Image
|
| import config
|
|
|
| def preprocess_image(image, target_size=512):
|
| """
|
| Preprocess the input image: resize and ensure it's RGB.
|
|
|
| Parameters:
|
| -----------
|
| image : PIL.Image or np.ndarray
|
| Input image
|
| target_size : int
|
| Target size for the image (will be square)
|
|
|
| Returns:
|
| --------
|
| np.ndarray : Preprocessed image
|
| """
|
| if isinstance(image, np.ndarray):
|
| image = Image.fromarray(image)
|
|
|
|
|
| if image.mode != 'RGB':
|
| image = image.convert('RGB')
|
|
|
|
|
| image = image.resize((target_size, target_size))
|
|
|
| return np.array(image)
|
|
|
| def divide_into_grid(image):
|
| """
|
| Divide image into a grid of cells.
|
|
|
| Parameters:
|
| -----------
|
| image : np.ndarray
|
| Input image
|
|
|
| Returns:
|
| --------
|
| list : List of cell information (row, col, cell_array, avg_color)
|
| """
|
| height, width = image.shape[:2]
|
| grid_h, grid_w = config.grid_size
|
|
|
| cell_height = height // grid_h
|
| cell_width = width // grid_w
|
|
|
| cells = []
|
|
|
| try :
|
| cells = image.reshape(grid_h, cell_height, grid_w, cell_width, 3)
|
| except Exception:
|
| raise TypeError("Image too big! The Grid doesn't divide the resized image evenly due to cascading rounding")
|
| cells = cells.transpose(0, 2, 1, 3, 4)
|
| cells = cells.reshape(-1, cell_height, cell_width, 3)
|
|
|
|
|
| vectorized = np.vectorize(compute_average_color, signature='(i,j,k)->(n)')
|
| cells = vectorized(cells)
|
| return cells
|
|
|
| def compute_average_color(cell):
|
| """
|
| Compute the average color of a cell.
|
|
|
| Parameters:
|
| -----------
|
| cell : np.ndarray
|
| Image cell
|
|
|
| Returns:
|
| --------
|
| np.ndarray : Average RGB color
|
| """
|
|
|
| return np.mean(cell, axis=(0,1)) |
