|
|
from transformers import ProcessorMixin |
|
|
from PIL import Image |
|
|
import torch |
|
|
import torchvision.transforms as T |
|
|
import numpy as np |
|
|
import os |
|
|
import json |
|
|
|
|
|
class DFineProcessor(ProcessorMixin): |
|
|
processor_class = "DFineProcessor" |
|
|
|
|
|
def __init__(self, size=640): |
|
|
self.size = size |
|
|
|
|
|
def resize_with_aspect_ratio(self, image, size): |
|
|
orig_w, orig_h = image.size |
|
|
ratio = min(size / orig_w, size / orig_h) |
|
|
new_w, new_h = int(orig_w * ratio), int(orig_h * ratio) |
|
|
image = image.resize((new_w, new_h), Image.BILINEAR) |
|
|
|
|
|
new_image = Image.new("RGB", (size, size)) |
|
|
pad_w, pad_h = (size - new_w) // 2, (size - new_h) // 2 |
|
|
new_image.paste(image, (pad_w, pad_h)) |
|
|
return new_image, ratio, pad_w, pad_h |
|
|
|
|
|
def __call__(self, images, return_tensors="pt"): |
|
|
if not isinstance(images, list): |
|
|
images = [images] |
|
|
|
|
|
processed_images = [] |
|
|
ratios = [] |
|
|
pad_ws = [] |
|
|
pad_hs = [] |
|
|
|
|
|
for image in images: |
|
|
if isinstance(image, np.ndarray): |
|
|
image = Image.fromarray(image[..., ::-1]) if image.shape[-1] == 3 else Image.fromarray(image) |
|
|
|
|
|
if not isinstance(image, Image.Image): |
|
|
raise ValueError("All inputs must be PIL images.") |
|
|
resized_img, ratio, pad_w, pad_h = self.resize_with_aspect_ratio(image, self.size) |
|
|
tensor_img = T.ToTensor()(resized_img) |
|
|
processed_images.append(tensor_img) |
|
|
ratios.append(ratio) |
|
|
pad_ws.append(pad_w) |
|
|
pad_hs.append(pad_h) |
|
|
|
|
|
torch_imgs = torch.stack(processed_images) |
|
|
ratios = torch.tensor(ratios) |
|
|
pad_w = torch.tensor(pad_ws) |
|
|
pad_h = torch.tensor(pad_hs) |
|
|
orig_target_sizes = torch.tensor([[self.size, self.size]]) |
|
|
|
|
|
return { |
|
|
"images": torch_imgs, |
|
|
"orig_target_sizes": orig_target_sizes, |
|
|
"ratio": ratios, |
|
|
"pad_w": pad_w, |
|
|
"pad_h": pad_h, |
|
|
} |
|
|
|
|
|
def save_pretrained(self, save_directory): |
|
|
os.makedirs(save_directory, exist_ok=True) |
|
|
with open(os.path.join(save_directory, "preprocessor_config.json"), "w") as f: |
|
|
json.dump({"processor_class": self.__class__.__name__}, f) |
|
|
|
|
|
@classmethod |
|
|
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): |
|
|
return cls() |
|
|
|
