Spaces:
Runtime error
Runtime error
File size: 1,501 Bytes
4a285f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
from typing import List
import numpy as np
from PIL import Image
from PIL.Image import Image as PILImage
from scipy.special import log_softmax
from .session_base import BaseSession
pallete1 = [
0,
0,
0,
255,
255,
255,
0,
0,
0,
0,
0,
0,
]
pallete2 = [
0,
0,
0,
0,
0,
0,
255,
255,
255,
0,
0,
0,
]
pallete3 = [
0,
0,
0,
0,
0,
0,
0,
0,
0,
255,
255,
255,
]
class ClothSession(BaseSession):
def predict(self, img: PILImage) -> List[PILImage]:
ort_outs = self.inner_session.run(
None, self.normalize(img, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), (768, 768))
)
pred = ort_outs
pred = log_softmax(pred[0], 1)
pred = np.argmax(pred, axis=1, keepdims=True)
pred = np.squeeze(pred, 0)
pred = np.squeeze(pred, 0)
mask = Image.fromarray(pred.astype("uint8"), mode="L")
mask = mask.resize(img.size, Image.LANCZOS)
masks = []
mask1 = mask.copy()
mask1.putpalette(pallete1)
mask1 = mask1.convert("RGB").convert("L")
masks.append(mask1)
mask2 = mask.copy()
mask2.putpalette(pallete2)
mask2 = mask2.convert("RGB").convert("L")
masks.append(mask2)
mask3 = mask.copy()
mask3.putpalette(pallete3)
mask3 = mask3.convert("RGB").convert("L")
masks.append(mask3)
return masks
|