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()
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()
My desired image should be like this
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!