File size: 2,063 Bytes
d972a8d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
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)
#combined_image = None
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 # To be modified
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))
|