|
|
import cv2, os, json |
|
|
import numpy as np |
|
|
from math import atan2, cos, sin, sqrt, pi |
|
|
from sklearn.cluster import DBSCAN |
|
|
from sklearn.preprocessing import StandardScaler |
|
|
from sklearn.cluster import OPTICS, cluster_optics_dbscan |
|
|
from tqdm import tqdm |
|
|
from glob import glob |
|
|
|
|
|
def nms(boxes, thresh): |
|
|
if len(boxes) == 0: |
|
|
return [] |
|
|
|
|
|
pick = [] |
|
|
x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] |
|
|
area = (x2 - x1 + 1) * (y2 - y1 + 1) |
|
|
idxs = np.argsort(y2) |
|
|
|
|
|
while len(idxs) > 0: |
|
|
last = len(idxs) - 1 |
|
|
i = idxs[last] |
|
|
pick.append(i) |
|
|
|
|
|
xx1 = np.maximum(x1[i], x1[idxs[:last]]) |
|
|
yy1 = np.maximum(y1[i], y1[idxs[:last]]) |
|
|
xx2 = np.minimum(x2[i], x2[idxs[:last]]) |
|
|
yy2 = np.minimum(y2[i], y2[idxs[:last]]) |
|
|
|
|
|
w = np.maximum(0, xx2 - xx1 + 1) |
|
|
h = np.maximum(0, yy2 - yy1 + 1) |
|
|
|
|
|
overlap = (w * h) / area[idxs[:last]] |
|
|
idxs = np.delete(idxs, np.concatenate(([last], np.where(overlap > thresh)[0]))) |
|
|
|
|
|
return boxes[pick] |
|
|
|
|
|
def find_pts_on_line(og, slope, d): |
|
|
cx, cy = og |
|
|
x1 = float(cx - d / ((1 + slope**2)**0.5)) |
|
|
y1 = float(cy - slope * cx + x1 * slope) |
|
|
return (x1, y1) |
|
|
|
|
|
def cluster(img_path, im, save_dir): |
|
|
img = cv2.imread(img_path) |
|
|
imgray = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) |
|
|
|
|
|
ret, binary_map = cv2.threshold(imgray, 127, 255, 0) |
|
|
|
|
|
nlabels, labels, stats, centroids = cv2.connectedComponentsWithStats(binary_map, None, None, None, 8, cv2.CV_32S) |
|
|
areas = stats[1:, cv2.CC_STAT_AREA] |
|
|
result = np.zeros((labels.shape), np.uint8) |
|
|
for i in range(0, nlabels - 1): |
|
|
if areas[i] >= 250: |
|
|
result[labels == i + 1] = 255 |
|
|
re_copy = result.copy() |
|
|
|
|
|
edgeimg = cv2.Canny(result, 10, 150) |
|
|
|
|
|
size = np.size(result) |
|
|
skel = np.zeros(result.shape, np.uint8) |
|
|
element = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3)) |
|
|
|
|
|
while True: |
|
|
open_ = cv2.morphologyEx(result, cv2.MORPH_OPEN, element) |
|
|
temp = cv2.subtract(result, open_) |
|
|
eroded = cv2.erode(result, element) |
|
|
skel = cv2.bitwise_or(skel, temp) |
|
|
result = eroded.copy() |
|
|
if cv2.countNonZero(result) == 0: |
|
|
break |
|
|
|
|
|
nlabels, labels, stats, centroids = cv2.connectedComponentsWithStats(skel, None, None, None, 8, cv2.CV_32S) |
|
|
areas = stats[1:, cv2.CC_STAT_AREA] |
|
|
skel = np.zeros((labels.shape), np.uint8) |
|
|
for i in range(0, nlabels - 1): |
|
|
if areas[i] >= 2: |
|
|
skel[labels == i + 1] = 255 |
|
|
|
|
|
|
|
|
base_name = os.path.splitext(im)[0] |
|
|
os.makedirs(save_dir, exist_ok=True) |
|
|
cv2.imwrite(os.path.join(save_dir, f"Skeleton_{base_name}.png"), skel) |
|
|
|
|
|
combined_img = cv2.addWeighted(skel, 1, edgeimg, 1, 0) |
|
|
cv2.imwrite(os.path.join(save_dir, f"SkeletonContour_{base_name}.png"), combined_img) |
|
|
|
|
|
filter_size = (10, 10) |
|
|
white_pixels = np.where(skel == 255) |
|
|
x_coords, y_coords = white_pixels[1], white_pixels[0] |
|
|
|
|
|
x1 = x_coords - filter_size[0] // 2 |
|
|
y1 = y_coords - filter_size[1] // 2 |
|
|
x2 = x_coords + filter_size[0] // 2 |
|
|
y2 = y_coords + filter_size[1] // 2 |
|
|
|
|
|
white_regions = np.column_stack((x1, y1, x2, y2)) |
|
|
white_regions = nms(white_regions, thresh=0.1) |
|
|
|
|
|
skeleton_image = cv2.cvtColor(skel.copy(), cv2.COLOR_GRAY2BGR) |
|
|
filtered_image = cv2.cvtColor(re_copy.copy(), cv2.COLOR_GRAY2BGR) |
|
|
center_points, directions = [], [] |
|
|
|
|
|
def get_direction2(bbox_pixels): |
|
|
nonzero_indices = np.column_stack(np.nonzero(bbox_pixels)) |
|
|
nonzero_indices = np.float32(nonzero_indices) |
|
|
if len(nonzero_indices) >= 2: |
|
|
mean, eigenvectors = cv2.PCACompute(nonzero_indices, mean=None) |
|
|
cntr = ((mean[0, 1]), (mean[0, 0])) |
|
|
return eigenvectors[0], cntr |
|
|
else: |
|
|
return (0, 0), (0, 0) |
|
|
|
|
|
for coor in white_regions: |
|
|
x1, y1, x2, y2 = coor |
|
|
bbox_pixels = skel[int(y1):int(y2), int(x1):int(x2)] |
|
|
direction, mean = get_direction2(bbox_pixels) |
|
|
directions.append(direction) |
|
|
center_points.append((mean[0] + x1, mean[1] + y1)) |
|
|
cv2.rectangle(skeleton_image, (int(x1), int(y1)), (int(x2), int(y2)), (255, 255, 255), 1) |
|
|
cv2.circle(skeleton_image, (int(mean[0] + x1), int(mean[1] + y1)), 1, (255, 0, 0), -1) |
|
|
|
|
|
cv2.imwrite(os.path.join(save_dir, f"Result_{base_name}.png"), skeleton_image) |
|
|
|
|
|
pts_group, bbox_group = [], [] |
|
|
for idx, pts in enumerate(center_points): |
|
|
if 640 > pts[0] > 0 and 480 > pts[1] > 0: |
|
|
pts_group.append([int(pts[0]), int(pts[1])]) |
|
|
x1, y1, x2, y2 = white_regions[idx] |
|
|
bbox_group.append([int(x1), int(y1), int(x2), int(y2)]) |
|
|
|
|
|
return pts_group, bbox_group |
|
|
|
|
|
|
|
|
mask_dir = "datasets/seg_train" |
|
|
save_json_dir = "datasets" |
|
|
save_img_dir = os.path.join(save_json_dir, "output") |
|
|
|
|
|
patterns = ['*.png', '*.jpg', '*.jpeg', '*.PNG', '*.JPG', '*.JPEG'] |
|
|
files = [] |
|
|
for p in patterns: |
|
|
files.extend(glob(os.path.join(mask_dir, p))) |
|
|
files = sorted(set(files)) |
|
|
print(f"Found {len(files)} files in {mask_dir}") |
|
|
|
|
|
file_dict = {} |
|
|
bbox_dict = {} |
|
|
|
|
|
for filepath in tqdm(files): |
|
|
filename = os.path.basename(filepath) |
|
|
pts, bbox = cluster(filepath, filename, save_img_dir) |
|
|
if len(pts) != 0: |
|
|
file_dict[filename] = pts |
|
|
bbox_dict[filename] = bbox |
|
|
|
|
|
with open(os.path.join(save_json_dir, 'train_seg_points.json'), 'w') as json_file: |
|
|
json.dump(file_dict, json_file) |
|
|
|
|
|
with open(os.path.join(save_json_dir, 'train_bbox_points.json'), 'w') as json_file: |
|
|
json.dump(bbox_dict, json_file) |
|
|
|
