CODEWITHSUNDEEP
pythontensorflowopencvimage-processingcrop
LuQchor
LuQchor

Reputation: 13

How do I crop an image based on custom mask

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

Upvotes: 1

Views: 1299

Answers (2)

asdf
asdf

Reputation: 1050

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)
    ymin, ymax = np.where(rows)[0][[0, -1]]
    xmin, xmax = np.where(cols)[0][[0, -1]]
    return ymin, ymax, xmin, xmax

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

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

Upvotes: 1

PepeChuy
PepeChuy

Reputation: 84

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]

Upvotes: 1

Related Questions