Home > database >  How do I crop an image based on custom mask
How do I crop an image based on custom mask

Time:09-01

I start my model and get a mask prediction of the area I would like to crop.

img = cv2.imread('picture.jpg')
img = cv2.resize(img, (224, 224))
print(img.shape)
T = np.zeros((1, 224, 224, 3), dtype='float32')
T[0]=img
prediction = model.predict(T , verbose=1)
prediction = prediction[0, :, : , :]
plt.imshow(img)
plt.show()
print(prediction.shape)
plt.imshow(np.squeeze(prediction).astype(np.float32))
plt.show()

enter image description here

This code is the closest I have gotten to the desired image, the problem with this image that is provided is that it doesn't crop out black background

result = prediction * img
plt.imshow((result/255).astype(np.float32))
plt.show()

enter image description here

My desired image should be like this enter image description here

CodePudding user response:

You could try this:

prediction = np.squeeze(prediction).astype(np.float32)
threshold = 0.5
xmin,ymin,xmax,ymax = (-1,-1,-1,-1)
for j in range(224):
    for i in range(224):
        if prediction[j,i] <= threshold: continue
        xmin = i if xmin == -1 or xmin > i else xmin
        xmax = i if xmax == -1 or xmax < i else xmax
        ymin = j if ymin == -1 or ymin > j else ymin
        ymax = j if ymax == -1 or ymax < j else ymax
cropImg = img[ymin:ymax,xmin:xmax]

CodePudding user response:

You can try to find the bounding box around your predicted mask and then crop your image using the detected bounding box. Using the code from this excellent SO answer, it would result in something like this:

def bbox2(img):
    # from https://stackoverflow.com/a/31402351/19249364
    rows = np.any(img, axis=1)
    cols = np.any(img, axis=0)
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]
    return rmin, rmax, cmin, cmax

bbox = bbox2(prediction.squeeze())
cropped = img[bbox[0]:bbox[1] 1, bbox[2]:bbox[3] 1, :]

Couldn't test your code, just tell me if it doesn't work!

  • Related