41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
|
import numpy as np
|
||
|
from PIL import Image
|
||
|
import sys
|
||
|
|
||
|
def preprocess(mask_values, pil_img, scale, is_mask):
|
||
|
w, h = pil_img.size
|
||
|
newW, newH = int(scale * w), int(scale * h)
|
||
|
assert newW > 0 and newH > 0, 'Scale is too small, resized images would have no pixel'
|
||
|
pil_img = pil_img.resize((newW, newH), resample=Image.NEAREST if is_mask else Image.BICUBIC)
|
||
|
img = np.asarray(pil_img)
|
||
|
|
||
|
if is_mask:
|
||
|
mask = np.zeros((newH, newW), dtype=np.int64)
|
||
|
for i, v in enumerate(mask_values):
|
||
|
if img.ndim == 2:
|
||
|
mask[img == v] = i
|
||
|
else:
|
||
|
mask[(img == v).all(-1)] = i
|
||
|
|
||
|
return mask
|
||
|
|
||
|
else:
|
||
|
if img.ndim == 2:
|
||
|
img = img[np.newaxis, ...]
|
||
|
else:
|
||
|
img = img.transpose((2, 0, 1))
|
||
|
|
||
|
if (img > 1).any():
|
||
|
img = img / 255.0
|
||
|
|
||
|
return img
|
||
|
|
||
|
if __name__=='__main__':
|
||
|
img = Image.open(sys.argv[1])
|
||
|
img = preprocess(None, img, 0.1, False)
|
||
|
img = img.astype(np.float32)
|
||
|
print(img.shape)
|
||
|
|
||
|
with open(sys.argv[2],'wb') as f:
|
||
|
f.write(img.tobytes())
|