-
Notifications
You must be signed in to change notification settings - Fork 56
/
Copy pathpreprocessing.py
204 lines (167 loc) · 8.76 KB
/
preprocessing.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
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
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
import json
import os
import cv2
import random
import os
import pandas as pd
import numpy as np
from detectron2.structures import BoxMode
from tqdm import tqdm
def get_image_ids(image_folder=None):
"""
Explores a folder of images and gets their ID from their file name.
Returns a list of all image ID's in image_folder.
E.g. image_folder/608fda8c976e0ac.jpg -> ["608fda8c976e0ac"]
Params
------
image_folder (str): path to folder of images, e.g. "../validation/"
"""
return [os.path.splitext(img_name)[0] for img_name in os.listdir(image_folder) if img_name.endswith(".jpg")]
# Make a function which formats a specific annotations csv based on what we're dealing with
def format_annotations(image_folder, annotation_file, target_classes=None):
"""
TODO - NOTE: This function could (definitely can) be faster.
TODO - Some ideas: skip the use of pandas entirely and use CSV's
Formats annotation_file based on images contained in image_folder.
Will get all unique image IDs and make sure annotation_file
only contains those (the target images).
Adds meta-data to annotation_file such as class names and categories.
If target_classes isn't None, the returned annotations will be filtered by this list.
Note: image_folder and annotation_file should both be validation if working on
validation set or both be training if working on training set.
Params
------
image_folder (str): path to folder of target images.
annotation_file (str): path to annotation file of target images.
target_classes (list), optional: a list of target classes you'd like to filter labels.
"""
# Get all image ids from target directory
image_ids = get_image_ids(image_folder)
# Setup annotation file and classnames
# TODO - improve this, is pandas required?
annot_file = pd.read_csv(annotation_file)
classes = pd.read_csv("class-descriptions-boxable.csv",
names=["LabelName", "ClassName"])
# Create classname column on annotations which converts label codes to string labels
annot_file["ClassName"] = annot_file["LabelName"].map(classes.set_index("LabelName")["ClassName"])
# Sort annot_file by "ClassName" for alphabetical labels (used with target_classes)
annot_file.sort_values(by=["ClassName"], inplace=True)
# TODO - fix this, Make sure we only get the images we're concerned about
if target_classes:
annot_file = annot_file[annot_file["ImageID"].isin(image_ids) & annot_file["ClassName"].isin(target_classes)]
else:
annot_file = annot_file[annot_file["ImageID"].isin(image_ids)]
assert len(annot_file.ImageID.unique()) == len(image_ids), "Label unique ImageIDs doesn't match target folder."
# Add ClassID column, e.g. "Bathtub, Toilet" -> 1, 2
annot_file["ClassName"] = pd.Categorical(annot_file["ClassName"])
annot_file["ClassID"] = annot_file["ClassName"].cat.codes
return annot_file
def rel_to_absolute(bbox, height, width):
"""
Converts bounding box dimensions from relative to absolute pixel values (Detectron2 style).
See: https://detectron2.readthedocs.io/modules/structures.html#detectron2.structures.BoxMode
Params
------
bbox (array): relative dimensions of bounding box in format (x0, y0, x1, y1 or Xmin, Ymin, Xmax, Ymax)
height (int): height of image
width (int): width of image
"""
bbox[0] = np.round(np.multiply(bbox[0], width)) # x0
bbox[1] = np.round(np.multiply(bbox[1], height)) # y0
bbox[2] = np.round(np.multiply(bbox[2], width)) # x1
bbox[3] = np.round(np.multiply(bbox[3], height)) # y1
return [i.astype("object") for i in bbox] # convert all to objects for JSON saving
def get_image_dicts(image_folder, annotation_file, target_classes=None):
"""
Create JSON of dectectron2 style labels to be reused later.
TODO -- Maybe create some verbosity here? AKA, what are the outputs?
TODO -- what if annotations = None? Can we create a call to create an annotations CSV in 1 hit?
Params
------
image_folder (str): target folder containing images
annotations (DataFrame): DataFrame of image label data
"""
# Get name of dataset from image_folder
dataset_name = str(image_folder)
print(f"Using {annotation_file} for annotations...")
# TODO: there should be some kind of asssertions here making sure the image folder and annotation files match
# E.g. train w/ train and valid w/ valid
annotations = format_annotations(image_folder=image_folder,
annotation_file=annotation_file,
target_classes=target_classes)
print(f"On dataset: {dataset_name}")
print("Classes we're using:\n {}".format(annotations["ClassName"].value_counts()))
# Get all unique image ids from target folder
img_ids = get_image_ids(image_folder)
print(f"Total number of images: {len(img_ids)}")
# TODO: move img_data creation out of for loop and only work with subset of img_ids?
#img_data = annotations[annotations["ImageID"] == img].reset_index() # reset index important for images with multiple objects
#change to something like "img_data = annotations is in img_ids..."
# Start creating image dictionaries (Detectron2 style labelling)
img_dicts = []
for idx, img in tqdm(enumerate(img_ids)):
record = {}
# Get image metadata
file_name = image_folder + "/" + img + ".jpg"
height, width = cv2.imread(file_name).shape[:2]
img_data = annotations[annotations["ImageID"] == img].reset_index() # reset index important for images
# with multiple objects
# Verbosity for image label troubleshooting
# print(f"On image: {img}")
# print(f"Image category: {img_data.ClassID.values}")
# print(f"Image label: {img_data.ClassName.values}")
# Update record dictionary
record["file_name"] = file_name
record["image_id"] = idx
record["height"] = height
record["width"] = width
# Create list of image annotations (labels)
img_annotations = []
for i in range(len(img_data)): # this is where we loop through examples with multiple objects in an image
category_id = img_data.loc[i]["ClassID"].astype("object") # JSON (for evalution) can't take int8 (NumPy type) must be native Python type
# print(f"Image category 2: {category_id}")
# Get bounding box coordinates in Detectron2 style (x0, y0, x1, y1)
bbox = np.float32(img_data.loc[i][["XMin", "YMin", "XMax", "YMax"]].values) # needs to be float/int # TODO: change for JSON
# Convert bbox from relative to absolute pixel dimensions
bbox = rel_to_absolute(bbox=bbox, height=height, width=width)
# Setup annot (1 annot = 1 label, there might be more) dictionary
annot = {
"bbox": bbox,
"bbox_mode": BoxMode.XYXY_ABS, # See: https://detectron2.readthedocs.io/modules/structures.html#detectron2.structures.BoxMode.XYXY_ABS
"category_id": category_id
}
img_annotations.append(annot)
# Update record dictionary with annotations
record["annotations"] = img_annotations
# Add record dictionary with image annotations to img_dicts list
img_dicts.append(record)
# TODO: Change this into it's own function??
# Save img_dicts to JSON for use later
json_file = os.path.join(image_folder, dataset_name+"_labels.json")
print(f"Saving labels to: {json_file}...")
with open(json_file, "w") as f:
json.dump(img_dicts, f)
# return img labels dictionary
return img_dicts
def load_json_labels(image_folder):
"""
Returns Detectron2 style labels of images in image_folder based on JSON label file in image_folder.
TODO -- Maybe create some verbosity here? AKA, what are the outputs?
TODO -- what if annotations = None? Can we create a call to create an annotations CSV in 1 hit?
Params
------
image_folder (str): target folder containing images
"""
# Get absolute path of JSON label file
for file in os.listdir(image_folder):
if file.endswith(".json"):
json_file = os.path.join(image_folder, file)
# TODO: Fix this assertion
assert json_file, "No .json label file found, please make one with annots_to_json()"
with open(json_file, "r") as f:
img_dicts = json.load(f)
# Convert bbox_mode to Enum of BoxMode.XYXY_ABS (doesn't work loading normal from JSON)
for img_dict in img_dicts:
for annot in img_dict["annotations"]:
annot["bbox_mode"] = BoxMode.XYXY_ABS
return img_dicts