| import numpy as np |
| import cv2 |
| import cv2 as cv |
| from PIL import Image |
| import random |
| from torchvision.transforms import transforms |
|
|
| def compute(patch): |
| patch = np.array(patch).astype(np.int64) |
| diff_horizontal = np.sum(np.abs(patch[:, :-1, :] - patch[:, 1:, :])) |
| diff_vertical = np.sum(np.abs(patch[:-1, :, :] - patch[1:, :, :])) |
| diff_diagonal = np.sum(np.abs(patch[:-1, :-1, :] - patch[1:, 1:, :])) |
| diff_diagonal += np.sum(np.abs(patch[1:, :-1, :] - patch[:-1, 1:, :])) |
| res = diff_horizontal + diff_vertical + diff_diagonal |
| return res.sum() |
|
|
| def bit_patch(img, img_height, bit_mode, patch_size, patch_mode): |
| img_np = np.array(img) |
| |
| if bit_mode == "scaling": |
| mask_low = 0x07 |
| red_low3 = ((img_np[:, :, 0] & mask_low) * (255 // 7)).astype(np.uint8) |
| green_low3 = ((img_np[:, :, 1] & mask_low) * (255 // 7)).astype(np.uint8) |
| blue_low3 = ((img_np[:, :, 2] & mask_low) * (255 // 7)).astype(np.uint8) |
| combined_image = cv.merge((red_low3, green_low3, blue_low3)) |
| elif bit_mode == "thresholding": |
| combined_image=img_np |
| else: |
| raise ValueError(f"Unsupported bit_mode: {bit_mode}") |
| h, w, _ = combined_image.shape |
| combined_image = Image.fromarray(combined_image) |
| min_len = min(h, w) |
| rz = transforms.Resize((img_height, img_height)) |
| if min_len < patch_size: |
| combined_image = rz(combined_image) |
| num_patch = (img_height // patch_size) ** 2 |
| patch_list = [] |
| rp = transforms.RandomCrop(patch_size) |
| for _ in range(num_patch): |
| patch_list.append(rp(combined_image)) |
|
|
| if patch_mode == "max": |
| patch_list.sort(key=lambda x: compute(x), reverse=True) |
| selected_patch = patch_list[0] |
| elif patch_mode == "min": |
| patch_list.sort(key=lambda x: compute(x), reverse=False) |
| selected_patch = patch_list[0] |
| else: |
| selected_patch = random.choice(patch_list) |
|
|
| patch_np = np.array(selected_patch) |
| return cv2.resize(patch_np, (img_height, img_height)) |
|
|
