在python中使用相应的蒙版进行图像分割

Question

I have corresponding masks to the images that I want to segment.

I put the images in one folder and their corresponding masks in another folder. I'm trying to apply those masks or multiply them by the images using two for loops in python to get the segmented images.

I'm using the code below:

def ImageSegmentation():

SegmentedImages = []

for img_path in os.listdir('C:/Users/mab/Desktop/images/'):
    img=io.imread('C:/Users/mab/Desktop/data/'+img_path)
    for img_path2 in os.listdir('C:/Users/mab/Desktop/masks/'):
        Mask = io.imread('C:/Users/mab/Desktop/masks/'+img_path2)

        [indx, indy] = np.where(Mask==0)
        Color_Masked = img.copy()
        Color_Masked[indx,indy] = 0

        matplotlib.image.imsave('C:/Users/mab/Desktop/SegmentedImages/'+img_path2,Color_Masked)
        segs.append(Color_Masked)
 return np.vstack(Color_Masked)

This code works when I try it for a single image and a single mask (without the folders and loops).

However, when I try to loop over the images and masks I have in the two folders, I get output images that are segmented by the wrong mask (not their corresponding mask).

I can't segment each single image alone without looping because I have more than 500 Images and their masks.

I don't know what I'm missing or placing wrong in this code and how can I fix it? Also, is there an easier way to get the segmented images?

Answer 1

Unless I have grossly misunderstood, you just need something like this:

import glob

filelist = glob.glob('C:/Users/mab/Desktop/images/*.png') 
for i in filelist:
    mask = i.replace("images","masks")
    print(i,mask)

On my iMac, that sort of thing produces:

/Users/mark/StackOverflow/images/b.png /Users/mark/StackOverflow/masks/b.png
/Users/mark/StackOverflow/images/a.png /Users/mark/StackOverflow/masks/a.png