ImageSegmentation/landmark_code.py at master · micahzev/ImageSegmentation · 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

"""

implements a very useful Landmark object that always returns a 40x2 dimensional landmark object with points

allows for basic to complex landmark manipulations and analysis

"""


import os
import re
import numpy as np

class Landmark(object):

    def __init__(self, landmark_data):
        if isinstance(landmark_data, str):
            lines = open(landmark_data).readlines()
            points = []
            for x, y in zip(lines[0::2], lines[1::2]):
                points.append(np.array([float(x), float(y)]))
            self.points = np.array(points)

        elif isinstance(landmark_data, np.ndarray) and np.atleast_2d(landmark_data).shape[0] == 1:
            self.points = np.array((landmark_data[:len(landmark_data)/2], landmark_data[len(landmark_data)/2:])).T

        elif isinstance(landmark_data, np.ndarray) and landmark_data.shape[1] == 2:
            self.points = landmark_data


    def as_vector(self):
        return np.hstack((self.points[:, 0], self.points[:, 1]))

    def get_center(self):
        return [self.points[:, 0].min() + (self.points[:, 0].max() - self.points[:, 0].min())/2,
                self.points[:, 1].min() + (self.points[:, 1].max() - self.points[:, 1].min())/2]

    def get_crown(self, is_upper):
        if is_upper:
            return Landmark(self.points[10:30, :])
        else:
            points = np.vstack((self.points[0:10, :], self.points[30:40, :]))
            return Landmark(points)

    def translate_to_origin(self):
        centroid = np.mean(self.points, axis=0)
        points = self.points - centroid
        return Landmark(points)

    def scale_to_unit(self):
        centroid = np.mean(self.points, axis=0)
        scale_factor = np.sqrt(np.power(self.points - centroid, 2).sum())
        points = self.points.dot(1. / scale_factor)
        return Landmark(points)

    def translate(self, vec):
        points = self.points + vec
        return Landmark(points)

    def scale(self, factor):
        centroid = np.mean(self.points, axis=0)
        points = (self.points - centroid).dot(factor) + centroid
        return Landmark(points)

    def capture_bounding_box(self, bbox):
        bbox_h = bbox[1][1] - bbox[0][1]
        scale_h = bbox_h / (self.points[:, 1].max() - self.points[:, 1].min())
        return self.scale(scale_h)

    def scaleposition(self, factor):
        points = self.points.dot(factor)
        return Landmark(points)

    def rotate(self, angle):
        rotmat = np.array([[np.cos(angle), np.sin(angle)],[-np.sin(angle), np.cos(angle)]])

        points = np.zeros_like(self.points)
        centroid = np.mean(self.points, axis=0)
        tmp_points = self.points - centroid
        for ind in range(len(tmp_points)):
            points[ind, :] = tmp_points[ind, :].dot(rotmat)
        points = points + centroid

        return Landmark(points)

    def T(self, t, s, theta):
        return self.rotate(theta).scale(s).translate(t)

    def invT(self, t, s, theta):
        return self.translate(-t).scale(1/s).rotate(-theta)


def load_landmarks(landmarks_dir, incisor, mirrored, exclude=None):

    search_string = "-" + str(incisor)

    if exclude:
        original_landmark_files = [landmarks_dir + 'original/' + file for file in os.listdir(landmarks_dir + 'original')
                                   if search_string in file and "s" + str(exclude)+"-" not in file]
    else:
        original_landmark_files = [landmarks_dir + 'original/' + file for file in os.listdir(landmarks_dir + 'original')
                                   if search_string in file]

    files = sorted(original_landmark_files, key=lambda x: int(re.search('[0-9]+', x).group()))

    if mirrored:
        if exclude:
            mirrored_landmark_files = [landmarks_dir + 'mirrored/' + file for file in os.listdir(landmarks_dir + 'mirrored')
                                        if search_string in file and "s" + str(exclude+14)+"-" not in file]
        else:
            mirrored_landmark_files = [landmarks_dir + 'mirrored/' + file for file in os.listdir(landmarks_dir + 'mirrored')
                                       if search_string in file]

        files2 = sorted(mirrored_landmark_files, key=lambda x: int(re.search('[0-9]+', x).group()))
        files = files + files2

    landmarks = []
    for filename in files:
        landmarks.append(Landmark(filename))
    return landmarks


def load_ground(landmarks_dir, incisor, l_o_o):

    search_string = "-" + str(incisor)

    ground = [landmarks_dir + 'original/' + file for file in os.listdir(landmarks_dir + 'original')
                                   if search_string in file and "s"+str(l_o_o)+"-" in file]

    return Landmark(ground[0])


def load_all_landmarks(landmarks_dir, exclude=None):

    search_string = "-"

    if exclude:
        original_landmark_files = [landmarks_dir + 'original/' + file for file in os.listdir(landmarks_dir + 'original')
                                   if "s" + str(exclude)+"-" not in file]
    else:
        original_landmark_files = [landmarks_dir + 'original/' + file for file in os.listdir(landmarks_dir + 'original')
                                   if search_string in file]

    files = sorted(original_landmark_files, key=lambda x: int(re.search('[0-9]+', x).group()))


    landmarks = []
    for filename in files:
        landmarks.append(Landmark(filename))
    return landmarks