Home > Enterprise >  How to cluster PyTorch predictions
How to cluster PyTorch predictions

Time:10-28

I'm trying to find road lanes from road images and then make predictions out of the images. So far, I've trained a model that finds road lanes. But most of the predictions are scattered. I'm trying to cluster PyTorch predictions that we get from these road images. These dots are the predictions of model where the road lanes might be.

Predictions shape: [1, 1, 80, 120]

Here's the image of predictions:

enter image description here

Here's what I want to achieve (I edited the image, deleted the dots that are scattered):

enter image description here

As you can see, I deleted the dots (predictions) from the image. I want each dot to be clustered with each other. How can I achieve this? I tried KNN (K Nearest Neighbors) but it didn't work.

CodePudding user response:

If you only want to remove dots then you can try to use enter image description here

Code:

import cv2
import numpy as np

mask = cv2.imread('road_mask.jpg', cv2.IMREAD_GRAYSCALE)
mask = cv2.resize(mask, (120, 80))

mask = cv2.erode(mask, np.ones((2, 2)))
mask = cv2.dilate(mask, np.ones((3, 3)))
mask = ((mask > 10) * 255).astype(np.uint8)

cv2.imwrite("postprocessed_mask.png", mask)
  • Related