Laudando-Associates-LLC/d-fine

1from transformers import ProcessorMixin2from PIL import Image3import torch4import torchvision.transforms as T5import numpy as np6import os7import json8 9class DFineProcessor(ProcessorMixin):10    processor_class = "DFineProcessor"11 12    def __init__(self, size=640):13        self.size = size14 15    def resize_with_aspect_ratio(self, image, size):16        orig_w, orig_h = image.size17        ratio = min(size / orig_w, size / orig_h)18        new_w, new_h = int(orig_w * ratio), int(orig_h * ratio)19        image = image.resize((new_w, new_h), Image.BILINEAR)20 21        new_image = Image.new("RGB", (size, size))22        pad_w, pad_h = (size - new_w) // 2, (size - new_h) // 223        new_image.paste(image, (pad_w, pad_h))24        return new_image, ratio, pad_w, pad_h25 26    def __call__(self, images, return_tensors="pt"):27        if not isinstance(images, list):28            images = [images]29 30        processed_images = []31        ratios = []32        pad_ws = []33        pad_hs = []34 35        for image in images:36            if isinstance(image, np.ndarray):37                image = Image.fromarray(image[..., ::-1]) if image.shape[-1] == 3 else Image.fromarray(image)38 39            if not isinstance(image, Image.Image):40                raise ValueError("All inputs must be PIL images.")41            resized_img, ratio, pad_w, pad_h = self.resize_with_aspect_ratio(image, self.size)42            tensor_img = T.ToTensor()(resized_img)43            processed_images.append(tensor_img)44            ratios.append(ratio)45            pad_ws.append(pad_w)46            pad_hs.append(pad_h)47 48        torch_imgs = torch.stack(processed_images)49        ratios = torch.tensor(ratios)50        pad_w = torch.tensor(pad_ws)51        pad_h = torch.tensor(pad_hs)52        orig_target_sizes = torch.tensor([[self.size, self.size]])53 54        return {55            "images": torch_imgs,56            "orig_target_sizes": orig_target_sizes,57            "ratio": ratios,58            "pad_w": pad_w,59            "pad_h": pad_h,60        }61 62    def save_pretrained(self, save_directory):63        os.makedirs(save_directory, exist_ok=True)64        with open(os.path.join(save_directory, "preprocessor_config.json"), "w") as f:65            json.dump({"processor_class": self.__class__.__name__}, f)66 67    @classmethod68    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):69        return cls()70