| |
|
|
| import cv2 |
| import numpy as np |
| from torchvision.transforms import transforms |
| from torchvision.transforms.functional import to_tensor |
| from transformers import CLIPProcessor |
|
|
| from basicsr.utils import img2tensor |
|
|
|
|
| class AddCannyFreezeThreshold(object): |
|
|
| def __init__(self, low_threshold=100, high_threshold=200): |
| self.low_threshold = low_threshold |
| self.high_threshold = high_threshold |
|
|
| def __call__(self, sample): |
| |
| x = sample['jpg'] |
| img = cv2.cvtColor(np.array(x), cv2.COLOR_RGB2BGR) |
| canny = cv2.Canny(img, self.low_threshold, self.high_threshold)[..., None] |
| sample['canny'] = img2tensor(canny, bgr2rgb=True, float32=True) / 255. |
| sample['jpg'] = to_tensor(x) |
| return sample |
|
|
|
|
| class AddCannyRandomThreshold(object): |
|
|
| def __init__(self, low_threshold=100, high_threshold=200, shift_range=50): |
| self.low_threshold = low_threshold |
| self.high_threshold = high_threshold |
| self.threshold_prng = np.random.RandomState() |
| self.shift_range = shift_range |
|
|
| def __call__(self, sample): |
| |
| x = sample['jpg'] |
| img = cv2.cvtColor(np.array(x), cv2.COLOR_RGB2BGR) |
| low_threshold = self.low_threshold + self.threshold_prng.randint(-self.shift_range, self.shift_range) |
| high_threshold = self.high_threshold + self.threshold_prng.randint(-self.shift_range, self.shift_range) |
| canny = cv2.Canny(img, low_threshold, high_threshold)[..., None] |
| sample['canny'] = img2tensor(canny, bgr2rgb=True, float32=True) / 255. |
| sample['jpg'] = to_tensor(x) |
| return sample |
|
|
|
|
| class AddStyle(object): |
|
|
| def __init__(self, version): |
| self.processor = CLIPProcessor.from_pretrained(version) |
| self.pil_to_tensor = transforms.ToTensor() |
|
|
| def __call__(self, sample): |
| |
| x = sample['jpg'] |
| style = self.processor(images=x, return_tensors="pt")['pixel_values'][0] |
| sample['style'] = style |
| sample['jpg'] = to_tensor(x) |
| return sample |
|
|
