cnc/preprocessing.py at master · melrchen/cnc · GitHub

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
import pylab as pl
import numpy as np
import random
import matplotlib.pyplot as plt
import matplotlib.image as img
import cv2
import sklearn
import sklearn.cluster
import os

def color_dist(c1, c2):
    '''
    Helper function that defines the distance between two pixels.

    Args:
        ci (list)

    Returns:
        dist (int)
    '''
    dist = 0
    for element in c1-c2:
        dist += abs(element)
    return dist

#image=img.imread('doggo.png')
#imgplot = plt.imshow(image)
#plt.show()


def read_in(filename):
    '''
    Reads in image filename.

    Args:
        filename (str)

    Returns:
        bgr, yuv tuple of np.arrays
    '''

    # BGR Image
    bgr = cv2.imread(filename)
    # print("bgr size")
    # print(bgr.shape)
    #cv2.imshow('BGR',bgr)

    #YUV image
    yuv = cv2.cvtColor(bgr,cv2.COLOR_BGR2YUV)
    # print("yuv size")
    # print(yuv.shape)
    # cv2.imshow('YUV',yuv)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    return bgr, yuv


def discretize_colors(filename):
    '''
    Takes in filename and outputs discretized pixel color mapping.

    Args:
        filename (string)

    Returns:
        pixel_color_map (dict): Mapping of coordinate (tuple) to color (np array YUV)
    '''
    bgr, yuv = read_in(filename)
    pixel_color_map = {} # Maps tuple of pixel coordinate to [Y, U, V] vector
    pixel_color_array = np.zeros(yuv.shape)

    #KMeans to find colors
    w, h, d = original_shape = tuple(yuv.shape)
    assert d == 3
    image_array = np.reshape(yuv, (w * h, d))

    kmeans = sklearn.cluster.KMeans(n_clusters=8).fit(image_array)
    labels = kmeans.labels_
    cluster_centers = kmeans.cluster_centers_

    # print("labels")
    # print(labels)
    # print("cluster centers")
    # print(cluster_centers)

    """
    #picture of cluster colors
    cluster_pic = []
    for cluster in cluster_centers:
        for ii in range(0, 50):
            col = []
            for jj in range(0, 50):
                col.append(cluster)
            cluster_pic.append(np.array(col))
    cluster_pic = np.array(cluster_pic)
    cluster_pic = cluster_pic.astype(np.uint8)
    #print "CLUSTER SHAPE"
    #print cluster_pic.shape
    #print cluster_pic

    bgr_clusters = cv2.cvtColor(cluster_pic,cv2.COLOR_YUV2BGR)
    cv2.imshow('clusters', bgr_clusters)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    """

    for ii in range(0, yuv.shape[0]):
        for jj in range(0, yuv.shape[1]):
            pixel = (ii, jj)
            best = cluster_centers[0]
            best_dist = color_dist(yuv[ii][jj], cluster_centers[0])

            for color in cluster_centers:
                dist = color_dist(yuv[ii][jj], color)
                if best_dist > dist:
                    best_dist = dist
                    best = color
            pixel_color_map[pixel] = best
            pixel_color_array[ii][jj] = best

    return pixel_color_array


    # bgr_map = cv2.cvtColor(pixel_color_array.astype(np.uint8),cv2.COLOR_YUV2BGR)
    # cv2.imshow('lol', bgr_map)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()


if __name__ == "__main__":
    # 0-255?

    path = os.path.join(os.getcwd(), 'city.jpeg')
    yuv = read_in(path)[1]

    print(yuv.shape)
    y, u, v = cv2.split(yuv)
    # y = yuv[:,:,0]
    # print(y)
    # cv2.imshow('y', y)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    # print(np.amin(u), np.amax(u))
    # print(np.amin(v), np.amax(v))
    # print(u[120])
    # print(u[1])
    print(u.shape)
    print(type(u[0][0]))
    cv2.imshow('u', u)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    print(v)
    cv2.imshow('y', v)
    cv2.waitKey(0)
    cv2.destroyAllWindows()