utils/utils.py

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)