openai / glide-text2im

GLIDE: a diffusion-based text-conditional image synthesis model
MIT License
3.54k stars 500 forks source link

How could I load a mask generated by myself? #41

Closed LZRZJUT closed 2 years ago

LZRZJUT commented 2 years ago

At first, I would like to say it's a amazing work. But when I try to change the code of 'inpaint.py' for using my own mask dataset, I realized it is uneasy. Because the mask was set by three lines. So, I want to ask for the code to use my own mask dataset like generated by PCov. Thanks a lot.

woctezuma commented 2 years ago
LZRZJUT commented 2 years ago

Thank you very much for your help, but in fact I had already read this issue before I asked this question. Just now I've changed the code and it works, and I'll put up my code as well.

def resize(img, height, width, centerCrop=True):
    imgh, imgw = img.shape[0:2]

    if centerCrop and imgh != imgw:
        # center crop
        side = np.minimum(imgh, imgw)
        j = (imgh - side) // 2
        i = (imgw - side) // 2
        img = img[j:j + side, i:i + side, ...]

    img = np.array(Image.fromarray(img).resize((height,width)))

    return img

def load_mask(imgh, imgw):
    mask = imread('./image/in/2409_mask.png')   # mask must be 255 for hole in this InpaintingModel
    mask = resize(mask, imgh, imgw, centerCrop=False)
    #mask = rgb2gray(mask)
    mask = (mask > 0).astype(np.uint8) * 255       # threshold due to interpolation
    mask = 255 - mask
    return to_tensor(mask)

def to_tensor(img):
    img = Image.fromarray(img)
    img_t = torchvision.transforms.functional.to_tensor(img).float()
    return img_t

source_mask_64 = load_mask(64, 64)
source_mask_256 = load_mask(256, 256)

With the above code you can load a mask image of your own, but of course you can also modify the function to make it more convenient。 Thanks again! Have a nice day!