-
Notifications
You must be signed in to change notification settings - Fork 2
/
make_patterns.py
90 lines (73 loc) · 3.65 KB
/
make_patterns.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
import requests
from skimage.color import rgb2lab, rgb2hsv
import csv
from PIL import Image
import os
import pickle
import matplotlib.pyplot as plt
import time
import datetime
from math import cos, sin
import numpy as np
from segment import segment_image, get_color_groups, hex2rgb, enclosure_strengths
from thing import ColorGroup, Pattern, ColorGroupSegment
folder = 'good'
def preprocess():
id2palette = {}
with open('good_dataset.csv') as gd:
reader = csv.DictReader(gd)
for row in reader:
id2palette[row['id']] = row['palette']
pickle_file = '{}_patterns.pickle'.format(folder)
count = 0
with open(pickle_file, 'wb') as pf:
with open('{}_dataset2.csv'.format(folder), 'w') as datafile:
images = os.listdir(folder)
for im in images:
img_num = im.split('.')[0]
if not img_num in id2palette:
res = requests.get("http://www.colourlovers.com/api/pattern/" + img_num, params={"format": "json"})
while res.status_code != 200:
res = requests.get("http://www.colourlovers.com/api/pattern/" + img_num, params={"format": "json"})
p = res.json()[0]
# print(p)
palette = p['colors']
else:
palette = id2palette[img_num].split()
# print(palette)
with Image.open(os.path.join(folder, im)) as imgfile:
img = imgfile.convert('RGBA')
img = np.array(img)
# print('Creating pattern')
print(count, ':', img_num)
count += 1
height,width,d = img.shape
# print('Segmenting')
start = time.time()
print(' Start:', datetime.datetime.fromtimestamp(start).strftime('%H:%M:%S'))
segments,px2id,adjacency = segment_image(img, palette)
print(' Adj:', len(adjacency))
if len(adjacency) > 5000:
print(' TOO MANY SEGMENTS')
continue
# print('Counting enclosure_strengths')
enc_str = enclosure_strengths(px2id, len(adjacency), adjacency)
endtime = time.time()
print(' End:', datetime.datetime.fromtimestamp(endtime).strftime('%H:%M:%S'))
if endtime - start > 240: # 240 seconds = 4 min
print(' OVER TIME LIMIT:', endtime - start)
continue
new_palette = []
for color in palette:
if len(segments[color]) == 0:
segments.pop(color, None)
else:
new_palette.append(color)
patt = Pattern(img_num, width, height, segments, px2id, enc_str, new_palette, 0)
# print('Dumping')
pickle.dump(patt, pf, protocol=4)
# metadata = p["userName"] + ',' + str(p['id']) + ',' + p['imageUrl'] + ',' + ' '.join(new_palette) + '\n'
metadata = "{},{}\n".format(img_num, ' '.join(new_palette))
datafile.write(metadata)
if __name__ == "__main__":
preprocess()