Newer
Older
import os
import numpy as np
import json
from PIL import Image
from patchify import patchify
import tensorflow as tf
# # Example of how to use the class
# data_loader = DataLoader('/path/to/AR_Maumee.tif', '/path/to/AR_Maumee.json', patch_size=(256, 256, 3), overlap=30)
# processed_data = data_loader.get_processed_data()
# reconstructed_image = data_loader.reconstruct_data(processed_data['map_patches'])
class DataLoader:
"""
DataLoader class to load and process TIFF images and corresponding JSON files.
Attributes:
tiff_path (str): Path to the input TIFF image file.
json_path (str): Path to the input JSON file.
patch_size (tuple of int): Size of the patches to extract from the TIFF image.
overlap (int): Overlapping pixels between patches.
map_img (numpy.ndarray): Loaded and normalized map image.
orig_size (tuple of int): Original size of the map image.
json_data (dict): Loaded JSON data.
processed_data (dict): Processed data including map patches and legends.
Methods:
load_and_process(): Loads and processes the TIFF image and JSON data.
load_tiff(): Loads and normalizes the TIFF image.
load_json(): Loads JSON data.
process_legends(label_suffix, resize_to=(256, 256)): Processes legends based on label suffix.
process_data(): Extracts and processes data from the loaded TIFF and JSON.
get_processed_data(): Returns the processed data.
reconstruct_data(self, patches): unpatchify the prediction data.
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
"""
def __init__(self, tiff_path, json_path, patch_size=(256, 256, 3), overlap=30):
"""
Initializes DataLoader with specified file paths, patch size, and overlap.
Parameters:
tiff_path (str): Path to the input TIFF image file.
json_path (str): Path to the input JSON file.
patch_size (tuple of int, optional): Size of patches to extract from image. Default is (256, 256, 3).
overlap (int, optional): Number of overlapping pixels between patches. Default is 30.
"""
self.tiff_path = tiff_path
self.json_path = json_path
self.patch_size = patch_size
self.overlap = overlap
self.map_img = None
self.orig_size = None
self.json_data = None
self.processed_data = None
self.load_and_process()
def load_and_process(self):
"""Loads and processes the TIFF image and JSON data."""
self.load_tiff()
self.load_json()
self.process_data()
def load_tiff(self):
"""Loads and normalizes the TIFF image."""
print(f"Loading and normalizing map image: {self.tiff_path}")
self.map_img = Image.open(self.tiff_path)
self.orig_size = self.map_img.size
self.map_img = np.array(self.map_img) / 255.0
def load_json(self):
"""Loads JSON data."""
with open(self.json_path, 'r') as json_file:
self.json_data = json.load(json_file)
def process_legends(self, label_suffix, resize_to=(256, 256)):
"""
Processes legends based on the label suffix.
Parameters:
label_suffix (str): Suffix in the label to identify the type of legend.
resize_to (tuple of int, optional): Size to resize the legends to. Default is (256, 256).
Returns:
list of tuples: List of processed legends and their corresponding labels.
"""
legends = [shape for shape in self.json_data['shapes'] if label_suffix in shape['label']]
processed_legends = []
for legend in legends:
points = np.array(legend['points'])
top_left = points.min(axis=0)
bottom_right = points.max(axis=0)
legend_img = self.map_img[int(top_left[1]):int(bottom_right[1]), int(top_left[0]):int(bottom_right[0]), :]
legend_img = tf.image.resize(legend_img, resize_to)
processed_legends.append((legend_img.numpy(), legend['label']))
return processed_legends
def process_data(self):
"""Extracts and processes data from the loaded TIFF and JSON."""
step_size = self.patch_size[0] - self.overlap
pad_x = (step_size - (self.map_img.shape[1] % step_size)) % step_size
pad_y = (step_size - (self.map_img.shape[0] % step_size)) % step_size
self.map_img = np.pad(self.map_img, ((0, pad_y), (0, pad_x), (0, 0)), mode='constant')
print(f"Patchifying map image with overlap...")
map_patches = patchify(self.map_img, self.patch_size, step=step_size)
poly_legends = self.process_legends('_poly')
pt_legends = self.process_legends('_pt')
line_legends = self.process_legends('_line')
self.processed_data = {
"map_patches": map_patches,
"poly_legends": poly_legends,
"pt_legends": pt_legends,
"line_legends": line_legends,
"original_size": self.orig_size,
}
def get_processed_data(self):
"""
Returns the processed data.
Raises:
ValueError: If data has not been loaded and processed.
Returns:
dict: Processed data including map patches and legends.
"""
if not self.processed_data:
raise ValueError("Data should be loaded and processed first")
return self.processed_data
def reconstruct_data(self, patches):
"""
Reconstructs an image from overlapping patches, keeping the maximum value
for overlapping pixels.
Parameters:
patches (numpy array): Image patches.
Returns:
numpy array: Reconstructed image.
"""
assert self.overlap >= 0, "Overlap should be non-negative"
step = patches.shape[2] - self.overlap
img_shape = self.orig_size[::-1] # Reverse the original size tuple to match the array shape
img = np.zeros(img_shape) # Initialize the image with zeros
for i in range(patches.shape[0]):
for j in range(patches.shape[1]):
x_start = i * step
y_start = j * step
x_end = min(x_start + patches.shape[2], img_shape[0])
y_end = min(y_start + patches.shape[3], img_shape[1])
img[x_start:x_end, y_start:y_end] = np.maximum(
img[x_start:x_end, y_start:y_end],
patches[i, j, :x_end-x_start, :y_end-y_start]
)
return img