-
Notifications
You must be signed in to change notification settings - Fork 2
/
labelbox_to_pickle.py
111 lines (77 loc) · 3.61 KB
/
labelbox_to_pickle.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
from collections import defaultdict
import pandas as pd
import sys
from skimage import io
import json
import cv2
import numpy as np
import pickle
"""Script to convert LabelBox CSV to Pandas DF"""
pd.set_option('display.max_colwidth', 1000)
input_csv_file = sys.argv[1]
output_dir = sys.argv[2]
dataset_csv = sys.argv[3]
def process_labels(labels):
if labels.get('colors_used', 'single_color') == 'single_color':
colors_used = 1
else:
# More than one color used
colors_used = 0
return labels.get('graffiti_type', 'tag'), \
colors_used, \
labels.get('readable_text', ''), \
labels.get('tool_used', 'marker')
def process_masks(labels, image_shape):
if 'Background Graffiti' in labels['segmentationMasksByName']:
background_graffiti = cv2.bitwise_not(cv2.inRange(io.imread(labels['segmentationMasksByName']['Background Graffiti']), 0, 0))
else:
background_graffiti = np.zeros(image_shape[0:2], dtype=np.uint8)
return cv2.bitwise_not(cv2.inRange(io.imread(labels['segmentationMasksByName']['Background']), 0, 0)),\
cv2.bitwise_not(cv2.inRange(io.imread(labels['segmentationMasksByName']['Graffiti']), 0, 0)), \
cv2.bitwise_not(cv2.inRange(io.imread(labels['segmentationMasksByName']['Incomplete Graffiti']), 0, 0)), \
background_graffiti
def check_pickle_file(pickle_file):
sample = pickle.load(open(pickle_file, 'rb'))
cv2.imwrite(f'test_images/{sample["id"]}_image.png', np.array(sample['image']))
cv2.imwrite(f'test_images/{sample["id"]}_background_mask.png', sample['background_mask'])
cv2.imwrite(f'test_images/{sample["id"]}_graffiti_mask.png', sample['graffiti_mask'])
cv2.imwrite(f'test_images/{sample["id"]}_incomplete_graffiti.png', sample['incomplete_graffiti_mask'])
cv2.imwrite(f'test_images/{sample["id"]}_background_graffiti_mask.png', sample['background_graffiti_mask'])
df = pd.DataFrame.from_csv(input_csv_file)
for col in list(df):
print(df.head()[col])
print()
original_datset_df = pd.read_csv(
dataset_csv,
index_col='hash_average',
usecols=['gps_latitude','gps_longitude','hash_average','original_file_name']
)
print(original_datset_df)
for index, row in df.iterrows():
data_sample = {}
try:
labels = json.loads(row['Label'])
except json.decoder.JSONDecodeError:
continue
image = io.imread(row['Labeled Data'])
print(labels)
graffiti_type, colors_used, readable_text, tool_used = process_labels(labels)
file_id = row['External ID'].split('.')[0]
print(file_id)
data_sample['id'] = file_id
data_sample['image'] = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
data_sample['graffiti_type'] = graffiti_type
data_sample['colors_used'] = colors_used
data_sample['readable_text'] = readable_text
data_sample['tool_used'] = tool_used
background_mask, graffiti_mask, incomplete_graffiti_mask, background_graffiti_mask = process_masks(labels, image.shape)
data_sample['background_mask'] = background_mask
data_sample['graffiti_mask'] = graffiti_mask
data_sample['incomplete_graffiti_mask'] = incomplete_graffiti_mask
data_sample['background_graffiti_mask'] = background_graffiti_mask
entry_from_original_datset = original_datset_df.loc[original_datset_df.index == file_id]
data_sample['gps_longitude'] = entry_from_original_datset['gps_longitude'].values[0]
data_sample['gps_latitude'] = entry_from_original_datset['gps_latitude'].values[0]
print(data_sample)
pickle.dump(data_sample, open(f'{output_dir}/{file_id}.p', 'wb'))
check_pickle_file(f'{output_dir}/{file_id}.p')