lama

1
import PIL.Image as Image
2
import numpy as np
3
import os
4

5

6
def main(args):
7
    if not args.indir.endswith('/'):
8
        args.indir += '/'
9
    os.makedirs(args.outdir, exist_ok=True)
10

11
    src_images = [
12
        args.indir+fname for fname in  os.listdir(args.indir)]
13

14
    tgt_masks = [
15
        args.outdir+fname[:-4] + f'_mask000.png' 
16
            for fname in  os.listdir(args.indir)]
17

18
    for img_name, msk_name in zip(src_images, tgt_masks):
19
        #print(img)
20
        #print(msk)
21

22
        image = Image.open(img_name).convert('RGB')
23
        image = np.transpose(np.array(image), (2, 0, 1))
24

25
        mask = (image == 255).astype(int)
26

27
        print(mask.dtype, mask.shape)
28

29

30
        Image.fromarray(
31
            np.clip(mask[0,:,:] * 255, 0, 255).astype('uint8'),mode='L'
32
        ).save(msk_name)
33

34

35

36

37
    '''
38
    for infile in src_images:
39
        try:
40
            file_relpath = infile[len(indir):]
41
            img_outpath = os.path.join(outdir, file_relpath)
42
            os.makedirs(os.path.dirname(img_outpath), exist_ok=True)
43

44
            image = Image.open(infile).convert('RGB')
45

46
            mask = 
47

48
            Image.fromarray(
49
                np.clip(
50
                    cur_mask * 255, 0, 255).astype('uint8'),
51
                    mode='L'
52
                ).save(cur_basename + f'_mask{i:03d}.png')
53
    '''
54

55

56

57
if __name__ == '__main__':
58
    import argparse
59
    aparser = argparse.ArgumentParser()
60
    aparser.add_argument('--indir', type=str, help='Path to folder with images')
61
    aparser.add_argument('--outdir', type=str, help='Path to folder to store aligned images and masks to')
62

63
    main(aparser.parse_args())
64