thinh-huynh-re
commited on
Commit
·
a0a24e3
1
Parent(s):
dcacd5e
Upload processor
Browse files- image_processor_dptdepth.py +47 -0
- preprocessor_config.json +7 -0
image_processor_dptdepth.py
ADDED
@@ -0,0 +1,47 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Dict, Optional, Tuple
|
2 |
+
|
3 |
+
import numpy as np
|
4 |
+
import torch.nn.functional as F
|
5 |
+
import torchvision.transforms as transforms
|
6 |
+
import torchvision.transforms.functional as TF
|
7 |
+
from PIL.Image import Image
|
8 |
+
from torch import Tensor
|
9 |
+
from transformers.image_processing_utils import BaseImageProcessor
|
10 |
+
|
11 |
+
INPUT_IMAGE_SIZE = (352, 352)
|
12 |
+
|
13 |
+
transform = transforms.Compose(
|
14 |
+
[
|
15 |
+
transforms.Resize(
|
16 |
+
(INPUT_IMAGE_SIZE, INPUT_IMAGE_SIZE),
|
17 |
+
interpolation=TF.InterpolationMode.BICUBIC,
|
18 |
+
),
|
19 |
+
transforms.Normalize(
|
20 |
+
(0.5, 0.5, 0.5),
|
21 |
+
(0.5, 0.5, 0.5),
|
22 |
+
),
|
23 |
+
]
|
24 |
+
)
|
25 |
+
|
26 |
+
|
27 |
+
class DPTDepthImageProcessor(BaseImageProcessor):
|
28 |
+
model_input_names = ["dptdepth_preprocessor"]
|
29 |
+
|
30 |
+
def __init__(self, testsize: Optional[int] = 352, **kwargs) -> None:
|
31 |
+
super().__init__(**kwargs)
|
32 |
+
self.testsize = testsize
|
33 |
+
|
34 |
+
def preprocess(
|
35 |
+
self, inputs: Dict[str, Image], **kwargs # {'rgb': ... }
|
36 |
+
) -> Dict[str, Tensor]:
|
37 |
+
return dict(rgb=transform(inputs["rgb"]).unsqueeze(0))
|
38 |
+
|
39 |
+
def postprocess(
|
40 |
+
self, logits: Tensor, size: Tuple[int, int], **kwargs
|
41 |
+
) -> np.ndarray:
|
42 |
+
logits: Tensor = F.upsample(
|
43 |
+
logits, size=size, mode="bilinear", align_corners=False
|
44 |
+
)
|
45 |
+
res: np.ndarray = logits.squeeze().data.cpu().numpy()
|
46 |
+
# res = (res - res.min()) / (res.max() - res.min() + 1e-8)
|
47 |
+
return res
|
preprocessor_config.json
ADDED
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"auto_map": {
|
3 |
+
"AutoImageProcessor": "image_processor_dptdepth.DPTDepthImageProcessor"
|
4 |
+
},
|
5 |
+
"image_processor_type": "DPTDepthImageProcessor",
|
6 |
+
"testsize": 352
|
7 |
+
}
|