-
Notifications
You must be signed in to change notification settings - Fork 0
/
superpixel_image_creation.py
137 lines (104 loc) · 4.91 KB
/
superpixel_image_creation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
from pycocotools.coco import COCO
import numpy as np
import cv2
from skimage.segmentation import slic
from skimage.future import graph
# from skimage.future import show_rag
import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import os
import json
# Path to the data
data_dir = './Data/'
annotation_file_training = data_dir + 'stuff_train2017.json'
annotation_file_val = data_dir + 'stuff_val2017.json'
image_dir_training = data_dir + 'train2017/'
image_dir_val = data_dir + 'val2017/'
segmentation_dir = data_dir + 'segmentations/'
val_segmentation_dir = data_dir + 'val_segmentations/'
os.makedirs(segmentation_dir, exist_ok=True)
#pixel_maps_dir_training = data_dir + 'stuff_train2017_pixelmaps/'
# Custom JSONEncoder that converts NumPy types to Python types
class NumpyEncoder(json.JSONEncoder):
""" Custom encoder for numpy data types """
def default(self, obj):
if isinstance(obj, np.ndarray):
return obj.tolist()
elif isinstance(obj, np.integer):
return int(obj)
elif isinstance(obj, np.floating):
return float(obj)
else:
return super(NumpyEncoder, self).default(obj)
def assign_labels_to_superpixels(annotations, segments, categories):
# Initialize a dictionary to hold the category label for each superpixel
superpixel_labels = {seg_id: np.zeros(len(categories)) for seg_id in np.unique(segments)}
# Go through each annotation and assign category labels to the superpixels
for ann in annotations:
# Get the category id, and find the corresponding index for the category
category_id = ann['category_id']
category_index = next((index for (index, c) in enumerate(categories) if c['id'] == category_id), None)
# Create a mask for the current annotation
# If the annotation is in polygon format, you will need to convert it to mask
ann_mask = coco.annToMask(ann) if ann['iscrowd'] == 0 else coco.decode(ann['segmentation'])
# For each superpixel, check if it overlaps with the annotation mask
for seg_id in superpixel_labels.keys():
superpixel_mask = segments == seg_id
overlap = ann_mask[superpixel_mask].sum()
# If there is significant overlap, mark the superpixel with the category label
if overlap > 0: # This threshold can be adjusted
superpixel_labels[seg_id][category_index] = 1
return superpixel_labels
# Initialize COCO api for instance annotations
coco = COCO(annotation_file_val)
# Load the categories
categories = coco.loadCats(coco.getCatIds())
category_names = [cat['name'] for cat in categories]
print('COCO categories: \n{}\n'.format(' '.join(category_names)))
# Get image ids
image_ids = coco.getImgIds()
# Load an image
for i, img_id in enumerate(image_ids[:3]):
# for img_id in image_ids:
image_info = coco.loadImgs(img_id)[0]
# image = cv2.imread(image_dir_training + image_info['file_name'])
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.imread(image_dir_val + image_info['file_name'])
if image is None:
continue # If the image can't be loaded, skip to the next one
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Apply superpixel segmentation to the image using SLIC
segments = slic(image, n_segments=250, compactness=10, sigma=1, start_label=1)
# Create graph based on superpixel adjacency
g = graph.rag_mean_color(image, segments)
# Convert to networkx graph to apply GNNs
nx_graph = nx.Graph()
for region in g:
for neighbor in g[region]:
edge_weight = g[region][neighbor]['weight'] # Example property
nx_graph.add_edge(region, neighbor, weight=edge_weight)
# Assuming g is your region adjacency graph from skimage
nx_graph = nx.Graph()
for region in g:
coords = np.column_stack(np.where(segments == region))
centroid = np.mean(coords, axis=0) # Calculate the centroid of the region
# Add each region as a node with a centroid attribute
nx_graph.add_node(region, centroid=centroid)
# Add edges with weights
for neighbor in g[region]:
edge_weight = g[region][neighbor]['weight']
nx_graph.add_edge(region, neighbor, weight=edge_weight)
# Create a plot with the original image
fig, ax = plt.subplots(figsize=(10, 10))
ax.imshow(image)
# For each node, find the centroid and draw
for region in nx_graph.nodes:
centroid = nx_graph.nodes[region]['centroid'] # Access the centroid attribute
ax.text(centroid[1], centroid[0], str(region), color='white', ha='center', va='center')
# Draw edges between regions
for (n1, n2) in nx_graph.edges:
p1 = nx_graph.nodes[n1]['centroid']
p2 = nx_graph.nodes[n2]['centroid']
ax.plot([p1[1], p2[1]], [p1[0], p2[0]], 'red') # Draw a line between centroids
plt.show()