creative-graphic-design/BASNet-SmartText
017
1from typing import Dict, Tuple, Union2 3import cv24import numpy as np5import torch6from PIL import Image7from PIL.Image import Image as PilImage8from torchvision import transforms9from transformers.image_processing_utils import BaseImageProcessor, BatchFeature10from transformers.image_utils import ImageInput11 12 13class RescaleT(object):14 def __init__(self, output_size: Union[int, Tuple[int, int]]) -> None:15 super().__init__()16 assert isinstance(output_size, (int, tuple))17 self.output_size = output_size18 19 def __call__(self, sample) -> Dict[str, np.ndarray]:20 image, label = sample["image"], sample["label"]21 22 h, w = image.shape[:2]23 24 if isinstance(self.output_size, int):25 if h > w:26 new_h, new_w = self.output_size * h / w, self.output_size27 else:28 new_h, new_w = self.output_size, self.output_size * w / h29 else:30 new_h, new_w = self.output_size31 32 new_h, new_w = int(new_h), int(new_w)33 34 # resize the image to new_h x new_w and convert image from range [0,255] to [0,1]35 # img = transform.resize(image,(new_h,new_w),mode='constant')36 # lbl = transform.resize(label,(new_h,new_w),mode='constant', order=0, preserve_range=True)37 38 # img = transform.resize(image, (self.output_size, self.output_size), mode='constant')39 img = (40 cv2.resize(41 image,42 (self.output_size, self.output_size),43 interpolation=cv2.INTER_AREA,44 )45 / 255.046 )47 # lbl = transform.resize(label, (self.output_size, self.output_size),48 # mode='constant',49 # order=0,50 # preserve_range=True)51 lbl = cv2.resize(52 label, (self.output_size, self.output_size), interpolation=cv2.INTER_NEAREST53 )54 lbl = np.expand_dims(lbl, axis=-1)55 lbl = np.clip(lbl, np.min(label), np.max(label))56 57 return {"image": img, "label": lbl}58 59 60class ToTensorLab(object):61 """Convert ndarrays in sample to Tensors."""62 63 def __init__(self, flag: int = 0) -> None:64 self.flag = flag65 66 def __call__(self, sample):67 image, label = sample["image"], sample["label"]68 69 tmpLbl = np.zeros(label.shape)70 71 if np.max(label) < 1e-6:72 label = label73 else:74 label = label / np.max(label)75 76 # change the color space77 if self.flag == 2: # with rgb and Lab colors78 tmpImg = np.zeros((image.shape[0], image.shape[1], 6))79 tmpImgt = np.zeros((image.shape[0], image.shape[1], 3))80 if image.shape[2] == 1:81 tmpImgt[:, :, 0] = image[:, :, 0]82 tmpImgt[:, :, 1] = image[:, :, 0]83 tmpImgt[:, :, 2] = image[:, :, 0]84 else:85 tmpImgt = image86 # tmpImgtl = color.rgb2lab(tmpImgt)87 tmpImgtl = cv2.cvtColor(tmpImgt, cv2.COLOR_RGB2LAB)88 89 # nomalize image to range [0,1]90 tmpImg[:, :, 0] = (tmpImgt[:, :, 0] - np.min(tmpImgt[:, :, 0])) / (91 np.max(tmpImgt[:, :, 0]) - np.min(tmpImgt[:, :, 0])92 )93 tmpImg[:, :, 1] = (tmpImgt[:, :, 1] - np.min(tmpImgt[:, :, 1])) / (94 np.max(tmpImgt[:, :, 1]) - np.min(tmpImgt[:, :, 1])95 )96 tmpImg[:, :, 2] = (tmpImgt[:, :, 2] - np.min(tmpImgt[:, :, 2])) / (97 np.max(tmpImgt[:, :, 2]) - np.min(tmpImgt[:, :, 2])98 )99 tmpImg[:, :, 3] = (tmpImgtl[:, :, 0] - np.min(tmpImgtl[:, :, 0])) / (100 np.max(tmpImgtl[:, :, 0]) - np.min(tmpImgtl[:, :, 0])101 )102 tmpImg[:, :, 4] = (tmpImgtl[:, :, 1] - np.min(tmpImgtl[:, :, 1])) / (103 np.max(tmpImgtl[:, :, 1]) - np.min(tmpImgtl[:, :, 1])104 )105 tmpImg[:, :, 5] = (tmpImgtl[:, :, 2] - np.min(tmpImgtl[:, :, 2])) / (106 np.max(tmpImgtl[:, :, 2]) - np.min(tmpImgtl[:, :, 2])107 )108 109 # tmpImg = tmpImg/(np.max(tmpImg)-np.min(tmpImg))110 111 tmpImg[:, :, 0] = (tmpImg[:, :, 0] - np.mean(tmpImg[:, :, 0])) / np.std(112 tmpImg[:, :, 0]113 )114 tmpImg[:, :, 1] = (tmpImg[:, :, 1] - np.mean(tmpImg[:, :, 1])) / np.std(115 tmpImg[:, :, 1]116 )117 tmpImg[:, :, 2] = (tmpImg[:, :, 2] - np.mean(tmpImg[:, :, 2])) / np.std(118 tmpImg[:, :, 2]119 )120 tmpImg[:, :, 3] = (tmpImg[:, :, 3] - np.mean(tmpImg[:, :, 3])) / np.std(121 tmpImg[:, :, 3]122 )123 tmpImg[:, :, 4] = (tmpImg[:, :, 4] - np.mean(tmpImg[:, :, 4])) / np.std(124 tmpImg[:, :, 4]125 )126 tmpImg[:, :, 5] = (tmpImg[:, :, 5] - np.mean(tmpImg[:, :, 5])) / np.std(127 tmpImg[:, :, 5]128 )129 130 elif self.flag == 1: # with Lab color131 tmpImg = np.zeros((image.shape[0], image.shape[1], 3))132 133 if image.shape[2] == 1:134 tmpImg[:, :, 0] = image[:, :, 0]135 tmpImg[:, :, 1] = image[:, :, 0]136 tmpImg[:, :, 2] = image[:, :, 0]137 else:138 tmpImg = image139 140 # tmpImg = color.rgb2lab(tmpImg)141 print("tmpImg:", tmpImg.min(), tmpImg.max())142 exit()143 tmpImg = cv2.cvtColor(tmpImg, cv2.COLOR_RGB2LAB)144 145 # tmpImg = tmpImg/(np.max(tmpImg)-np.min(tmpImg))146 147 tmpImg[:, :, 0] = (tmpImg[:, :, 0] - np.min(tmpImg[:, :, 0])) / (148 np.max(tmpImg[:, :, 0]) - np.min(tmpImg[:, :, 0])149 )150 tmpImg[:, :, 1] = (tmpImg[:, :, 1] - np.min(tmpImg[:, :, 1])) / (151 np.max(tmpImg[:, :, 1]) - np.min(tmpImg[:, :, 1])152 )153 tmpImg[:, :, 2] = (tmpImg[:, :, 2] - np.min(tmpImg[:, :, 2])) / (154 np.max(tmpImg[:, :, 2]) - np.min(tmpImg[:, :, 2])155 )156 157 tmpImg[:, :, 0] = (tmpImg[:, :, 0] - np.mean(tmpImg[:, :, 0])) / np.std(158 tmpImg[:, :, 0]159 )160 tmpImg[:, :, 1] = (tmpImg[:, :, 1] - np.mean(tmpImg[:, :, 1])) / np.std(161 tmpImg[:, :, 1]162 )163 tmpImg[:, :, 2] = (tmpImg[:, :, 2] - np.mean(tmpImg[:, :, 2])) / np.std(164 tmpImg[:, :, 2]165 )166 167 else: # with rgb color168 tmpImg = np.zeros((image.shape[0], image.shape[1], 3))169 image = image / np.max(image)170 if image.shape[2] == 1:171 tmpImg[:, :, 0] = (image[:, :, 0] - 0.485) / 0.229172 tmpImg[:, :, 1] = (image[:, :, 0] - 0.485) / 0.229173 tmpImg[:, :, 2] = (image[:, :, 0] - 0.485) / 0.229174 else:175 tmpImg[:, :, 0] = (image[:, :, 0] - 0.485) / 0.229176 tmpImg[:, :, 1] = (image[:, :, 1] - 0.456) / 0.224177 tmpImg[:, :, 2] = (image[:, :, 2] - 0.406) / 0.225178 179 tmpLbl[:, :, 0] = label[:, :, 0]180 181 # change the r,g,b to b,r,g from [0,255] to [0,1]182 # transforms.Normalize(mean = (0.485, 0.456, 0.406), std = (0.229, 0.224, 0.225))183 tmpImg = tmpImg.transpose((2, 0, 1))184 tmpLbl = label.transpose((2, 0, 1))185 186 return {"image": torch.from_numpy(tmpImg), "label": torch.from_numpy(tmpLbl)}187 188 189def apply_transform(190 data: Dict[str, np.ndarray], rescale_size: int, to_tensor_lab_flag: int191) -> Dict[str, torch.Tensor]:192 transform = transforms.Compose(193 [RescaleT(output_size=rescale_size), ToTensorLab(flag=to_tensor_lab_flag)]194 )195 return transform(data) # type: ignore196 197 198class BASNetImageProcessor(BaseImageProcessor):199 model_input_names = ["pixel_values"]200 201 def __init__(202 self, rescale_size: int = 256, to_tensor_lab_flag: int = 0, **kwargs203 ) -> None:204 super().__init__(**kwargs)205 self.rescale_size = rescale_size206 self.to_tensor_lab_flag = to_tensor_lab_flag207 208 def preprocess(self, images: ImageInput, **kwargs) -> BatchFeature:209 if not isinstance(images, PilImage):210 raise ValueError(f"Expected PIL.Image, got {type(images)}")211 212 image_pil = images213 image_npy = np.array(image_pil, dtype=np.uint8)214 width, height = image_pil.size215 label_npy = np.zeros((height, width), dtype=np.uint8)216 217 assert image_npy.shape[-1] == 3218 output = apply_transform(219 {"image": image_npy, "label": label_npy},220 rescale_size=self.rescale_size,221 to_tensor_lab_flag=self.to_tensor_lab_flag,222 )223 image = output["image"]224 225 assert isinstance(image, torch.Tensor)226 227 return BatchFeature(228 data={"pixel_values": image.float().unsqueeze(dim=0)}, tensor_type="pt"229 )230 231 def postprocess(232 self, prediction: torch.Tensor, width: int, height: int233 ) -> PilImage:234 def _norm_prediction(d: torch.Tensor) -> torch.Tensor:235 ma, mi = torch.max(d), torch.min(d)236 237 # division while avoiding zero division238 dn = (d - mi) / ((ma - mi) + torch.finfo(torch.float32).eps)239 return dn240 241 prediction = _norm_prediction(prediction)242 prediction = prediction.squeeze()243 prediction = prediction * 255 + 0.5244 prediction = prediction.clamp(0, 255)245 246 prediction_np = prediction.cpu().numpy()247 image = Image.fromarray(prediction_np).convert("RGB")248 image = image.resize((width, height), resample=Image.Resampling.BILINEAR)249 return image250