jagennath-hari commited on
Commit
5211df4
·
1 Parent(s): 2e571c4

add D-Fine preprocessor

Browse files
Files changed (2) hide show
  1. preprocessor_config.json +5 -2
  2. processor_dfine.py +69 -0
preprocessor_config.json CHANGED
@@ -1,3 +1,6 @@
1
  {
2
- "processor_class": "DFineProcessor"
3
- }
 
 
 
 
1
  {
2
+ "processor_class": "DFineProcessor",
3
+ "auto_map": {
4
+ "AutoProcessor": "d-fine--processor_dfine.DFineProcessor"
5
+ }
6
+ }
processor_dfine.py ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import ProcessorMixin
2
+ from PIL import Image
3
+ import torch
4
+ import torchvision.transforms as T
5
+ import numpy as np
6
+ import os
7
+ import json
8
+
9
+ class DFineProcessor(ProcessorMixin):
10
+ processor_class = "DFineProcessor"
11
+
12
+ def __init__(self, size=640):
13
+ self.size = size
14
+
15
+ def resize_with_aspect_ratio(self, image, size):
16
+ orig_w, orig_h = image.size
17
+ 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) // 2
23
+ new_image.paste(image, (pad_w, pad_h))
24
+ return new_image, ratio, pad_w, pad_h
25
+
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
+ @classmethod
68
+ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
69
+ return cls()