Spaces:
Running
Running
File size: 3,572 Bytes
3faa99b c8f8b0e 3faa99b c8f8b0e 3faa99b c8f8b0e 3faa99b c8f8b0e 3faa99b c8f8b0e 3faa99b c8f8b0e 3faa99b c8f8b0e 3faa99b c8f8b0e 3faa99b 5f57808 3faa99b c8f8b0e 3faa99b |
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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
import os
from typing import List
import numpy as np
import pooch
from PIL import Image
from PIL.Image import Image as PILImage
from scipy.special import log_softmax
from .base import BaseSession
palette1 = [
0,
0,
0,
255,
255,
255,
0,
0,
0,
0,
0,
0,
]
palette2 = [
0,
0,
0,
0,
0,
0,
255,
255,
255,
0,
0,
0,
]
palette3 = [
0,
0,
0,
0,
0,
0,
0,
0,
0,
255,
255,
255,
]
class Unet2ClothSession(BaseSession):
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Predict the cloth category of an image.
This method takes an image as input and predicts the cloth category of the image.
The method uses the inner_session to make predictions using a pre-trained model.
The predicted mask is then converted to an image and resized to match the size of the input image.
Depending on the cloth category specified in the method arguments, the method applies different color palettes to the mask and appends the resulting images to a list.
Parameters:
img (PILImage): The input image.
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
List[PILImage]: A list of images representing the predicted masks.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(
img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (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.Resampling.LANCZOS)
masks = []
cloth_category = kwargs.get("cc") or kwargs.get("cloth_category")
def upper_cloth():
mask1 = mask.copy()
mask1.putpalette(palette1)
mask1 = mask1.convert("RGB").convert("L")
masks.append(mask1)
def lower_cloth():
mask2 = mask.copy()
mask2.putpalette(palette2)
mask2 = mask2.convert("RGB").convert("L")
masks.append(mask2)
def full_cloth():
mask3 = mask.copy()
mask3.putpalette(palette3)
mask3 = mask3.convert("RGB").convert("L")
masks.append(mask3)
if cloth_category == "upper":
upper_cloth()
elif cloth_category == "lower":
lower_cloth()
elif cloth_category == "full":
full_cloth()
else:
upper_cloth()
lower_cloth()
full_cloth()
return masks
@classmethod
def download_models(cls, *args, **kwargs):
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:2434d1f3cb744e0e49386c906e5a08bb"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
return "u2net_cloth_seg"
|