Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import cv2 | |
| from collections import Counter | |
| COLOR_NAME = ['black', 'brown', 'blue', 'gray', 'green', 'orange', 'pink', 'purple', 'red', 'white', 'yellow'] | |
| def im2c(im, w2c): | |
| """ | |
| Convert an image to color name representation using a color-name matrix. | |
| Returns: | |
| numpy.ndarray: Processed image based on color parameter. | |
| """ | |
| # Define color name mappings | |
| color_values = np.array([[ 0, 0, 0], | |
| [165, 81, 43], | |
| [ 0, 0, 255], | |
| [127, 127, 127], | |
| [ 0, 255, 0], | |
| [255, 127, 0], | |
| [255, 165, 216], | |
| [191, 0, 191], | |
| [255, 0, 0], | |
| [255, 255, 255], | |
| [255, 255, 0]], dtype=np.uint8) | |
| # Extract RGB channels | |
| # RR, GG, BB = im[:, :, 0].flatten(), im[:, :, 1].flatten(), im[:, :, 2].flatten() | |
| # Compute index for w2c lookup | |
| index_im = ((im[:, :, 0].flatten() // 8) + 32 * (im[:, :, 1].flatten()// 8) + 32 * 32 * (im[:, :, 2].flatten() // 8)).astype(np.int32) | |
| # w2cM = np.argmax(w2c, axis=1) | |
| # name_idx_img = w2cM[index_im].reshape(im.shape[:2]) | |
| # max_prob = w2c[np.arange(w2c.shape[0]), w2cM] | |
| # max_prob_map = max_prob[index_im].reshape(im.shape[:2]) | |
| prob_map = w2c[index_im, :].reshape((im.shape[0], im.shape[1], w2c.shape[1])) | |
| max_prob_map = np.max(prob_map, axis=2) | |
| name_idx_img = np.argmax(prob_map, axis=2) | |
| color_img = np.zeros_like(im).astype(np.uint8) | |
| for jj in range(im.shape[0]): | |
| for ii in range(im.shape[1]): | |
| color_img[jj, ii, :] = np.array(color_values[name_idx_img[jj, ii]]) | |
| return prob_map, max_prob_map, name_idx_img, color_img | |
| if __name__ == "__main__": | |
| # Load w2c matrix | |
| w2c = np.load('w2c11_j.npy').astype(np.float16) # Assuming 'w2c.mat' was converted to 'w2c.npy' | |
| image_path = './test.jpg' | |
| img = cv2.imread(image_path).astype(np.float32) | |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| prob_map, max_prob_img, name_idx_img, color_img = im2c(img, w2c) | |
| filtered_counts = Counter(name_idx_img[name_idx_img <= 10]) | |
| sorted_counts = sorted(filtered_counts.items(), key=lambda x: x[1], reverse=True) | |
| top_3_values = [num for num, count in sorted_counts[:3]] | |
| top_3_colors = [COLOR_NAME[i] for i in top_3_values] | |
| print("Top 3 colors:", top_3_colors) | |
| cv2.imwrite('colormap_j.jpg', cv2.cvtColor(color_img, cv2.COLOR_BGR2RGB)) | |