image_processing_basnet.py

from typing import Dict, Tuple, Union

import cv2
import numpy as np
import torch
from PIL import Image
from PIL.Image import Image as PilImage
from torchvision import transforms
from transformers.image_processing_base import BatchFeature
from transformers.image_processing_utils import BaseImageProcessor
from transformers.image_utils import ImageInput


class RescaleT(object):
 def __init__(self, output_size: Union[int, Tuple[int, int]]) -> None:
 super().__init__()
 assert isinstance(output_size, (int, tuple))
 self.output_size = output_size

 def __call__(self, sample) -> Dict[str, np.ndarray]:
 image, label = sample["image"], sample["label"]

 h, w = image.shape[:2]

 if isinstance(self.output_size, int):
 if h > w:
 new_h, new_w = self.output_size * h / w, self.output_size
 else:
 new_h, new_w = self.output_size, self.output_size * w / h
 else:
 new_h, new_w = self.output_size

 new_h, new_w = int(new_h), int(new_w)

 # resize the image to new_h x new_w and convert image from range [0,255] to [0,1]
 # img = transform.resize(image,(new_h,new_w),mode='constant')
 # lbl = transform.resize(label,(new_h,new_w),mode='constant', order=0, preserve_range=True)

 # img = transform.resize(image, (self.output_size, self.output_size), mode='constant')
 img = (
 cv2.resize(
 image,
 (self.output_size, self.output_size),
 interpolation=cv2.INTER_AREA,
 )
 / 255.0
 )
 # lbl = transform.resize(label, (self.output_size, self.output_size),
 # mode='constant',
 # order=0,
 # preserve_range=True)
 lbl = cv2.resize(
 label, (self.output_size, self.output_size), interpolation=cv2.INTER_NEAREST
 )
 lbl = np.expand_dims(lbl, axis=-1)
 lbl = np.clip(lbl, np.min(label), np.max(label))

 return {"image": img, "label": lbl}


class ToTensorLab(object):
 """Convert ndarrays in sample to Tensors."""

 def __init__(self, flag: int = 0) -> None:
 self.flag = flag

 def __call__(self, sample):
 image, label = sample["image"], sample["label"]

 tmpLbl = np.zeros(label.shape)

 if np.max(label) < 1e-6:
 label = label
 else:
 label = label / np.max(label)

 # change the color space
 if self.flag == 2: # with rgb and Lab colors
 tmpImg = np.zeros((image.shape[0], image.shape[1], 6))
 tmpImgt = np.zeros((image.shape[0], image.shape[1], 3))
 if image.shape[2] == 1:
 tmpImgt[:, :, 0] = image[:, :, 0]
 tmpImgt[:, :, 1] = image[:, :, 0]
 tmpImgt[:, :, 2] = image[:, :, 0]
 else:
 tmpImgt = image
 # tmpImgtl = color.rgb2lab(tmpImgt)
 tmpImgtl = cv2.cvtColor(tmpImgt, cv2.COLOR_RGB2LAB)

 # nomalize image to range [0,1]
 tmpImg[:, :, 0] = (tmpImgt[:, :, 0] - np.min(tmpImgt[:, :, 0])) / (
 np.max(tmpImgt[:, :, 0]) - np.min(tmpImgt[:, :, 0])
 )
 tmpImg[:, :, 1] = (tmpImgt[:, :, 1] - np.min(tmpImgt[:, :, 1])) / (
 np.max(tmpImgt[:, :, 1]) - np.min(tmpImgt[:, :, 1])
 )
 tmpImg[:, :, 2] = (tmpImgt[:, :, 2] - np.min(tmpImgt[:, :, 2])) / (
 np.max(tmpImgt[:, :, 2]) - np.min(tmpImgt[:, :, 2])
 )
 tmpImg[:, :, 3] = (tmpImgtl[:, :, 0] - np.min(tmpImgtl[:, :, 0])) / (
 np.max(tmpImgtl[:, :, 0]) - np.min(tmpImgtl[:, :, 0])
 )
 tmpImg[:, :, 4] = (tmpImgtl[:, :, 1] - np.min(tmpImgtl[:, :, 1])) / (
 np.max(tmpImgtl[:, :, 1]) - np.min(tmpImgtl[:, :, 1])
 )
 tmpImg[:, :, 5] = (tmpImgtl[:, :, 2] - np.min(tmpImgtl[:, :, 2])) / (
 np.max(tmpImgtl[:, :, 2]) - np.min(tmpImgtl[:, :, 2])
 )

 # tmpImg = tmpImg/(np.max(tmpImg)-np.min(tmpImg))

 tmpImg[:, :, 0] = (tmpImg[:, :, 0] - np.mean(tmpImg[:, :, 0])) / np.std(
 tmpImg[:, :, 0]
 )
 tmpImg[:, :, 1] = (tmpImg[:, :, 1] - np.mean(tmpImg[:, :, 1])) / np.std(
 tmpImg[:, :, 1]
 )
 tmpImg[:, :, 2] = (tmpImg[:, :, 2] - np.mean(tmpImg[:, :, 2])) / np.std(
 tmpImg[:, :, 2]
 )
 tmpImg[:, :, 3] = (tmpImg[:, :, 3] - np.mean(tmpImg[:, :, 3])) / np.std(
 tmpImg[:, :, 3]
 )
 tmpImg[:, :, 4] = (tmpImg[:, :, 4] - np.mean(tmpImg[:, :, 4])) / np.std(
 tmpImg[:, :, 4]
 )
 tmpImg[:, :, 5] = (tmpImg[:, :, 5] - np.mean(tmpImg[:, :, 5])) / np.std(
 tmpImg[:, :, 5]
 )

 elif self.flag == 1: # with Lab color
 tmpImg = np.zeros((image.shape[0], image.shape[1], 3))

 if image.shape[2] == 1:
 tmpImg[:, :, 0] = image[:, :, 0]
 tmpImg[:, :, 1] = image[:, :, 0]
 tmpImg[:, :, 2] = image[:, :, 0]
 else:
 tmpImg = image

 # tmpImg = color.rgb2lab(tmpImg)
 print("tmpImg:", tmpImg.min(), tmpImg.max())
 exit()
 tmpImg = cv2.cvtColor(tmpImg, cv2.COLOR_RGB2LAB)

 # tmpImg = tmpImg/(np.max(tmpImg)-np.min(tmpImg))

 tmpImg[:, :, 0] = (tmpImg[:, :, 0] - np.min(tmpImg[:, :, 0])) / (
 np.max(tmpImg[:, :, 0]) - np.min(tmpImg[:, :, 0])
 )
 tmpImg[:, :, 1] = (tmpImg[:, :, 1] - np.min(tmpImg[:, :, 1])) / (
 np.max(tmpImg[:, :, 1]) - np.min(tmpImg[:, :, 1])
 )
 tmpImg[:, :, 2] = (tmpImg[:, :, 2] - np.min(tmpImg[:, :, 2])) / (
 np.max(tmpImg[:, :, 2]) - np.min(tmpImg[:, :, 2])
 )

 tmpImg[:, :, 0] = (tmpImg[:, :, 0] - np.mean(tmpImg[:, :, 0])) / np.std(
 tmpImg[:, :, 0]
 )
 tmpImg[:, :, 1] = (tmpImg[:, :, 1] - np.mean(tmpImg[:, :, 1])) / np.std(
 tmpImg[:, :, 1]
 )
 tmpImg[:, :, 2] = (tmpImg[:, :, 2] - np.mean(tmpImg[:, :, 2])) / np.std(
 tmpImg[:, :, 2]
 )

 else: # with rgb color
 tmpImg = np.zeros((image.shape[0], image.shape[1], 3))
 image = image / np.max(image)
 if image.shape[2] == 1:
 tmpImg[:, :, 0] = (image[:, :, 0] - 0.485) / 0.229
 tmpImg[:, :, 1] = (image[:, :, 0] - 0.485) / 0.229
 tmpImg[:, :, 2] = (image[:, :, 0] - 0.485) / 0.229
 else:
 tmpImg[:, :, 0] = (image[:, :, 0] - 0.485) / 0.229
 tmpImg[:, :, 1] = (image[:, :, 1] - 0.456) / 0.224
 tmpImg[:, :, 2] = (image[:, :, 2] - 0.406) / 0.225

 tmpLbl[:, :, 0] = label[:, :, 0]

 # change the r,g,b to b,r,g from [0,255] to [0,1]
 # transforms.Normalize(mean = (0.485, 0.456, 0.406), std = (0.229, 0.224, 0.225))
 tmpImg = tmpImg.transpose((2, 0, 1))
 tmpLbl = label.transpose((2, 0, 1))

 return {"image": torch.from_numpy(tmpImg), "label": torch.from_numpy(tmpLbl)}


def apply_transform(
 data: Dict[str, np.ndarray], rescale_size: int, to_tensor_lab_flag: int
) -> Dict[str, torch.Tensor]:
 transform = transforms.Compose(
 [RescaleT(output_size=rescale_size), ToTensorLab(flag=to_tensor_lab_flag)]
 )
 return transform(data) # type: ignore


class BASNetImageProcessor(BaseImageProcessor):
 model_input_names = ["pixel_values"]

 def __init__(
 self, rescale_size: int = 256, to_tensor_lab_flag: int = 0, **kwargs
 ) -> None:
 super().__init__(**kwargs)
 self.rescale_size = rescale_size
 self.to_tensor_lab_flag = to_tensor_lab_flag

 def preprocess(self, images: ImageInput, **kwargs) -> BatchFeature:
 if not isinstance(images, PilImage):
 raise ValueError(f"Expected PIL.Image, got {type(images)}")

 image_pil = images
 image_npy = np.array(image_pil, dtype=np.uint8)
 width, height = image_pil.size
 label_npy = np.zeros((height, width), dtype=np.uint8)

 assert image_npy.shape[-1] == 3
 output = apply_transform(
 {"image": image_npy, "label": label_npy},
 rescale_size=self.rescale_size,
 to_tensor_lab_flag=self.to_tensor_lab_flag,
 )
 image = output["image"]

 assert isinstance(image, torch.Tensor)

 return BatchFeature(
 data={"pixel_values": image.float().unsqueeze(dim=0)}, tensor_type="pt"
 )

 def postprocess(
 self, prediction: torch.Tensor, width: int, height: int
 ) -> PilImage:
 def _norm_prediction(d: torch.Tensor) -> torch.Tensor:
 ma, mi = torch.max(d), torch.min(d)

 # division while avoiding zero division
 dn = (d - mi) / ((ma - mi) + torch.finfo(torch.float32).eps)
 return dn

 # prediction = _norm_output(prediction)
 # prediction = prediction.squeeze()
 # prediction_np = prediction.cpu().numpy()

 # image = Image.fromarray(prediction_np * 255).convert("RGB")
 # image = image.resize((width, height), resample=Image.Resampling.BILINEAR)

 # return image

 # breakpoint()

 # output = F.interpolate(output, (height, width), mode="bilinear")
 # output = output.squeeze(dim=0)

 # output = _norm_output(output)

 # # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
 # output = output * 255 + 0.5
 # output = output.clamp(0, 255)

 # # shape: (C=1, W, H) -> (W, H, C=1)
 # output = output.permute(1, 2, 0)
 # # shape: (W, H, C=3)
 # output = output.repeat(1, 1, 3)

 # output_np = output.cpu().numpy().astype(np.uint8)
 # return Image.fromarray(output_np)

 prediction = _norm_prediction(prediction)
 prediction = prediction.squeeze()
 prediction = prediction * 255 + 0.5
 prediction = prediction.clamp(0, 255)

 prediction_np = prediction.cpu().numpy()
 image = Image.fromarray(prediction_np).convert("RGB")
 image = image.resize((width, height), resample=Image.Resampling.BILINEAR)
 return image