import cv2
from matplotlib import pyplot as plt
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
from datetime import datetime
import os
from typing import List, Dict
def convert_and_resize_mask(mask):
if mask.ndim == 3:
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
resized_mask = cv2.resize(mask, (1024, 1024))
return resized_mask
def add_masks_resized(masks):
final_mask = np.zeros((1024, 1024), dtype=np.uint8)
for mask in masks:
if mask is not None:
resized_mask = convert_and_resize_mask(mask)
resized_mask = resized_mask.astype(np.uint8)
final_mask = cv2.add(final_mask, resized_mask)
return final_mask
def attend_mask(mask_file, attend_scale=10, save=False):
if isinstance(mask_file, str):
if mask_file == '':
return torch.zeros([1, 1, 128, 128], dtype=torch.float32).cuda()
else:
image_with_mask = cv2.imread(mask_file, cv2.IMREAD_GRAYSCALE)
elif len(mask_file.shape) == 3: # convert RGB to gray
image_with_mask = cv2.cvtColor(mask_file, cv2.COLOR_BGR2GRAY)
else:
image_with_mask = mask_file
if attend_scale != 0:
kernel = np.ones((abs(attend_scale), abs(attend_scale)), np.uint8)
if attend_scale > 0:
image_with_mask = cv2.dilate(image_with_mask, kernel, iterations=1)
else:
image_with_mask = cv2.erode(image_with_mask, kernel, iterations=1)
if save and isinstance(mask_file, str):
new_mask_file_name = mask_file[:-4]+'_'+str(attend_scale)+'.jpg'
cv2.imwrite(new_mask_file_name, image_with_mask)
print("new_mask is saved in ", new_mask_file_name)
dilated_image= cv2.resize(image_with_mask, (128, 128), interpolation=cv2.INTER_NEAREST)
dilated_image = torch.from_numpy(dilated_image).to(torch.float32).unsqueeze(0).unsqueeze(0).cuda() / 255
return dilated_image
def panning(img_path=None, op_list=[['left', 0.2]], save=False, save_dir=None):
if isinstance(img_path, str):
img = cv2.imread(img_path)
else:
img = img_path
img_new = img.copy()
img_height, img_width, _ = img.shape
w_mask = 255 * np.ones((img_height, img_width), dtype=np.uint8)
h_mask = 255 * np.ones((img_height, img_width), dtype=np.uint8)
for op in op_list:
scale = op[1]
if op[0] in ['right', 'left']:
K = int(scale*img_width)
elif op[0] in ['up', 'down']:
K = int(scale*img_height)
if op[0] == 'right':
img_new[:, K:, :] = img[:, 0:img_width-K, :]
w_mask[:, K:] = 0
elif op[0] == 'left':
img_new[:, 0:img_width-K, :] = img[:, K:, :]
w_mask[:, 0:img_width-K] = 0
elif op[0] == 'down':
img_new[K:, :, :] = img[0:img_height-K, :, :]
h_mask[K:, :] = 0
elif op[0] == 'up':
img_new[0:img_height-K, :, :] = img[K:, :, :]
h_mask[0:img_height-K, :] = 0
img = img_new
mask = w_mask + h_mask
mask[mask>0] = 255
if save:
if save_dir is None:
base_dir = os.path.dirname(img_path)
save_dir = os.path.join(base_dir, 'preprocess')
elif not os.path.exists(save_dir):
os.makedirs(save_dir)
resized_img_name = f"{save_dir}/resized_image.png"
resized_mask_name = f"{save_dir}/resized_mask.png"
cv2.imwrite(resized_img_name, img_new)
cv2.imwrite(resized_mask_name, mask)
return resized_img_name, resized_mask_name
else:
return img_new, mask
def zooming(img_path=None, scale=[0.8, 0.8], save=False, save_dir=None):
if isinstance(img_path, str):
img = cv2.imread(img_path)
else:
img = img_path
img_new = img.copy()
img_height, img_width, _ = img.shape
mask = 255 * np.ones((img_height, img_width), dtype=np.uint8)
new_height = int(img_height*scale[0])
new_width = int(img_width*scale[1])
resized_img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)
x_offset = (img_width - new_width) // 2
y_offset = (img_height - new_height) // 2
img_new[y_offset:y_offset + new_height, x_offset:x_offset + new_width] = resized_img
mask[y_offset:y_offset + new_height, x_offset:x_offset + new_width] = 0
if save:
if save_dir is None:
base_dir = os.path.dirname(img_path)
save_dir = os.path.join(base_dir, 'preprocess')
elif not os.path.exists(save_dir):
os.makedirs(save_dir)
resized_img_name = f"{save_dir}/resized_image.png"
resized_mask_name = f"{save_dir}/resized_mask.png"
cv2.imwrite(resized_img_name, img_new)
cv2.imwrite(resized_mask_name, mask)
return resized_img_name, resized_mask_name
else:
return img_new, mask
def get_box(mask, bias = 2):
nonzero_indices = torch.nonzero(mask)
H, W = mask.shape[-2:]
min_x = max(min(nonzero_indices[:, 1]) - bias, 0)
min_y = max(min(nonzero_indices[:, 0]) - bias, 0)
max_x = min(max(nonzero_indices[:, 1]) + bias, W)
max_y = min(max(nonzero_indices[:, 0]) + bias, H)
return (min_x, min_y, max_x, max_y)
def draw_axis(img,grid_dict,x_len,y_len):
if grid_dict is not None and grid_dict is not False:
assert isinstance(grid_dict,Dict)
assert "x_title" in grid_dict
assert "y_title" in grid_dict
assert "x_text_list" in grid_dict
assert "y_text_list" in grid_dict
x_title=grid_dict["x_title"]
y_title=grid_dict["y_title"]
x_text_list=grid_dict['x_text_list']
y_text_list=grid_dict['y_text_list']
assert len(y_text_list)==y_len
assert len(x_text_list)==x_len
assert "font_size" in grid_dict
font_size=grid_dict["font_size"]
if "x_color" in grid_dict:
color_x=grid_dict['x_color']
else:
color_x="black"
if "y_color" in grid_dict:
color_y=grid_dict['y_color']
else:
color_y="black"
if "num_decimals" in grid_dict:
num_decimals=grid_dict['num_decimals']
else:
num_decimals=2
if "shift_x" in grid_dict:
shift_x_x,shift_x_y=grid_dict['shift_x']
else:
shift_x_x=shift_x_y=0
if "shift_y" in grid_dict:
shift_y_x,shift_y_y=grid_dict['shift_y']
else:
shift_y_x=shift_y_y=0
if "title" in grid_dict:
title=grid_dict['title']
if isinstance(title,List):
all_title=""
for s in title:
all_title=all_title+s+"\n"
title=all_title
else:
title=''
width, height = img.size
num_x=x_len
num_y=y_len
new_img = Image.new("RGB", (width + width // num_x+width // (num_x*2), height + height // num_y+height // (num_y*2)), color=(255, 255, 255))
width,height=(width + width // num_x, height + height // num_y)
num_x=num_x+1
num_y=num_y+1
new_img.paste(img, (width // num_x, height // num_y))
draw = ImageDraw.Draw(new_img)
font = ImageFont.truetype("DejaVuSansMono.ttf", font_size)
for i in range(2, num_x+1):
x = (i - 1) * width // num_x + width // (num_x * 2)-width *0.2// num_x+shift_x_x
y = height // (num_y * 2)+shift_x_y
k=i-1
if isinstance(x_text_list[i-2],str):
draw.text((x, y), x_text_list[i-2], font=font,fill=color_x,align="center")
else:
draw.text((x, y), "{:.{}f}".format(x_text_list[i-2],num_decimals), font=font,fill=color_x,align="center")
for i in range(2, num_y+1):
x = width // (num_x * 2)-width *0.1// num_x+shift_y_x
y = (i - 1) * height // num_y + height // (num_y * 2)-height*0.1//num_y+shift_y_y
k = i - 1
if isinstance(y_text_list[i-2],str):
draw.text((x, y), y_text_list[i-2], font=font,fill=color_y,align="center")
else:
draw.text((x, y), "{:.{}f}".format(y_text_list[i-2],num_decimals), font=font,fill=color_y,align="center")
i=1
x = (i - 1) * width // num_x + width // (num_x * 2)-height*0.1//num_y+shift_y_x
y = height // (num_y * 2)+width *0.2// num_x+shift_y_y
draw.text((x, y), y_title, font=font, fill=color_y,align="center")
x = width // (num_x * 2)+width *0.2// num_x+shift_x_x
y = (i - 1) * height // num_y + height // (num_y * 2)+shift_x_y
draw.text((x, y), x_title, font=font, fill=color_x,align="left")
x = width // 4
y = (i - 1) * height // num_y + height // (num_y * 10)
draw.text((x, y), title, font=font, fill='blue',align="left")
else:
new_img=img
return new_img
def view_images(images, num_rows=1, offset_ratio=0.02,text="",folder=None,Notimestamp=False,
grid_dict=None,subfolder=None,verbose=True,output_dir=None,timestamp=None,**kwargs):
if type(images) is list:
num_empty = len(images) % num_rows
elif images.ndim == 4:
num_empty = images.shape[0] % num_rows
else:
images = [images]
num_empty = 0
origin_size=kwargs.get("origin_size",None)
images_copy=images.copy()
for i, per_image in enumerate(images_copy):
if isinstance(per_image, Image.Image) and origin_size is not None:
images[i] = np.array(per_image.resize((origin_size[1],origin_size[0])))
else:
images[i] = np.array(per_image)
empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
num_items = len(images)
h, w, c = images[0].shape
offset = int(h * offset_ratio)
num_cols = num_items // num_rows
image_ = np.ones((h * num_rows + offset * (num_rows - 1),
w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
for i in range(num_rows):
for j in range(num_cols):
image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
i * num_cols + j]
pil_img = Image.fromarray(image_)
pil_img_=draw_axis(pil_img,grid_dict,num_cols,num_rows)
if pil_img_.size[0]==pil_img_.size[1]:
pil_img_.resize((2048,2048))
else:
longer_side = max(pil_img.size)
ratio = 2048/longer_side
new_size = tuple([int(x*ratio) for x in pil_img.size])
pil_img = pil_img.resize(new_size)
if verbose is False:
return pil_img
now = datetime.now()
if timestamp is None:
if Notimestamp is False:
timestamp = now.strftime("%Y-%m-%d_%H-%M-%S")
else:
timestamp=""
if output_dir is None:
if timestamp != "":
date, time = timestamp.split('_')
else:
date, time = "",""
if folder is not None:
dirname="./"+folder
filename = text+f"img_{timestamp}.jpg"
else:
if subfolder is not None:
dirname=os.path.join("./img", subfolder,date)
dirname=os.path.join(dirname,time)
filename =text+f"img_{timestamp}.jpg"
else:
dirname=os.path.join("./img",date)
dirname=os.path.join(dirname,time)
filename =text+f"img_{timestamp}.jpg"
else:
dirname=output_dir
filename =text+f"img_{timestamp}.jpg"
if not os.path.exists(dirname):
os.makedirs(dirname)
if verbose is True:
for i, img in enumerate(images):
im = Image.fromarray(img)
im.save(os.path.join(dirname,f"{i}.jpg"))
print(f"Output dir: {dirname}")
pil_img.save(os.path.join(dirname, filename))
if grid_dict is not None and grid_dict is not False:
if not os.path.exists(dirname):
os.makedirs(dirname)
pil_img_.save(os.path.join(dirname, filename[:-4]+"_2048x.jpg"))
def resize_image_with_mask(img, mask, scale):
if scale == 1:
return img, mask, None
img_blackboard = img.copy() # canvas
mask_blackboard = np.zeros_like(mask)
M = cv2.moments(mask)
cx = int(M["m10"] / M["m00"])
cy = int(M["m01"] / M["m00"])
scale_factor = [scale, scale]
resized_img = cv2.resize(img, None, fx=scale_factor[0], fy=scale_factor[1], interpolation=cv2.INTER_AREA)
resized_mask = cv2.resize(mask, None, fx=scale_factor[0], fy=scale_factor[1], interpolation=cv2.INTER_AREA)
new_cx, new_cy = cx * scale_factor[0], cy * scale_factor[1]
for y in range(resized_mask.shape[0]):
for x in range(resized_mask.shape[1]):
if 0 <= cy - (new_cy - y) < img.shape[0] and 0 <= cx - (new_cx - x) < img.shape[1]:
mask_blackboard[int(cy - (new_cy - y)), int(cx - (new_cx - x))] = resized_mask[y, x]
img_blackboard[int(cy - (new_cy - y)), int(cx - (new_cx - x))] = resized_img[y, x]
return img_blackboard, mask_blackboard, (cx, cy)
def flip_image_with_mask(img, mask, flip_code=None):
if flip_code is None:
return img, mask, None
M = cv2.moments(mask)
if M["m00"] == 0:
return img, mask
cx = int(M["m10"] / M["m00"])
cy = int(M["m01"] / M["m00"])
h, w = img.shape[:2]
img_center = (w // 2, h // 2)
tx = img_center[0] - cx
ty = img_center[1] - cy
M_translate = np.float32([[1, 0, tx], [0, 1, ty]])
img_translated = cv2.warpAffine(img, M_translate, (w, h))
mask_translated = cv2.warpAffine(mask, M_translate, (w, h))
flipped_img = cv2.flip(img_translated, flip_code)
flipped_mask = cv2.flip(mask_translated, flip_code)
M_translate_back = np.float32([[1, 0, -tx], [0, 1, -ty]])
flipped_img_back = cv2.warpAffine(flipped_img, M_translate_back, (w, h))
flipped_mask_back = cv2.warpAffine(flipped_mask, M_translate_back, (w, h))
return flipped_img_back, flipped_mask_back, (cx, cy)