MarkerAnnotation/sequence_model.py at master · MeteorsHub/MarkerAnnotation · 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
160
161
162
163
import os
import glob
from PIL import Image, ImageDraw
from collections import OrderedDict
import numpy as np
from functools import lru_cache


class SequenceModel:
    original_anno = [[-1, -1] for _ in range(10)]
    anno_radius = 3
    anno_color = [(255, 0, 0),
                  (255, 255, 0),
                  (0, 255, 0),
                  (0, 0, 255),
                  (255, 0, 255),

                  (255, 0, 0),
                  (255, 255, 0),
                  (0, 255, 0),
                  (0, 0, 255),
                  (255, 0, 255),
                  ]

    def __init__(self):
        self.working_dir = None
        self.annotation_file = None

        self.img_11_format = 'left_aligned_color_????.png'
        self.img_21_format = 'right_aligned_color_????.png'
        self.img_12_format = 'left_org_color_????.png'
        self.img_22_format = 'right_org_color_????.png'

        self.seq_list = OrderedDict()

    def clear_seq_list(self):
        self.seq_list.clear()

    @staticmethod
    def get_id_from_filename(filename):
        return int(filename[-8: -4])

    def load_working_dir(self, dir_name):
        if not os.path.exists(dir_name):
            raise FileExistsError('working_dir %s not found' % dir_name)
        try:
            self.save_annotation_file()
        except FileExistsError:
            pass
        self.clear_seq_list()
        self.working_dir = dir_name
        self.annotation_file = os.path.join(self.working_dir, 'annotation.txt')

        img_11_files = glob.glob(os.path.join(dir_name, self.img_11_format))
        seq_ids = [self.get_id_from_filename(file) for file in img_11_files]
        seq_ids.sort()

        seq_ids = [seq_ids[i] for i in range(0, len(seq_ids), 5)]

        annos = self.load_annotation_file()

        for seq_id in seq_ids:
            if seq_id in annos:
                anno = annos[seq_id]
            else:
                anno = None
            self.seq_list[seq_id] = anno
        return

    def load_annotation_file(self):
        anno = OrderedDict()
        if self.annotation_file is None or not os.path.exists(self.annotation_file):
            return anno
        with open(self.annotation_file) as f:
            for line in f.readlines():
                line = line.strip()
                if line and not line.startswith('#'):
                    items = line.split()
                    _id = int(items[0])
                    _anno_coords = []
                    for i in range(1, 21, 2):
                        _anno_coords.append([float(items[i]), float(items[i + 1])])
                    if _id in anno:
                        raise AttributeError('annotation file error')
                    anno[_id] = _anno_coords
        return anno

    def save_annotation_file(self, filename=None):
        if filename is None:
            if self.working_dir is None:
                return
            filename = self.annotation_file

        if filename is None:
            raise FileExistsError('saving annotation error for %s' % filename)
        with open(filename, 'w') as f:
            for _id, _anno in self.seq_list.items():
                if _anno is None:
                    continue
                f.write('%d\t' % _id)
                for i in range(10):
                    f.write('%1.2f\t%1.2f\t' % (_anno[i][0], _anno[i][1]))
                f.write('\n')
        return

    def update_anno(self, seq_id, finger_id, u, v):
        if seq_id not in self.seq_list:
            raise KeyError
        if self.seq_list[seq_id] is None:
            self.seq_list[seq_id] = self.original_anno.copy()
        self.seq_list[seq_id][finger_id] = [u, v]

    @lru_cache(10)
    def load_imgs(self, seq_id):
        if seq_id not in self.seq_list:
            return [None, None, None, None]
        img_files = [os.path.join(self.working_dir, self.img_11_format.replace('????', '%04d' % seq_id)),
                     os.path.join(self.working_dir, self.img_21_format.replace('????', '%04d' % seq_id)),
                     os.path.join(self.working_dir, self.img_12_format.replace('????', '%04d' % seq_id)),
                     os.path.join(self.working_dir, self.img_22_format.replace('????', '%04d' % seq_id))]
        imgs = []
        for file in img_files:
            if not os.path.exists(file):
                imgs.append(None)
            else:
                img = Image.open(file)
                img = img.resize((640, 480), Image.BILINEAR)
                img_array = np.asarray(img, np.uint8)
                imgs.append(img_array)
        return imgs

    def load_imgs_with_anno(self, seq_id):
        imgs = self.load_imgs(seq_id).copy()
        if self.seq_list[seq_id] is None:
            return imgs
        if imgs[0] is not None:
            img_11 = Image.fromarray(imgs[0])
            draw = ImageDraw.Draw(img_11)
            for i in range(5):
                anno = self.seq_list[seq_id][i]
                u = int(anno[0])
                v = int(anno[1])
                if u < 0 or v < 0:
                    continue
                draw.ellipse((u-self.anno_radius, v-self.anno_radius, u+self.anno_radius, v+self.anno_radius),
                             fill=self.anno_color[i])
            del draw
            imgs[0] = np.asarray(img_11, np.uint8)

        if imgs[1] is not None:
            img_21 = Image.fromarray(imgs[1])
            draw = ImageDraw.Draw(img_21)
            for i in range(5, 10):
                anno = self.seq_list[seq_id][i]
                u = int(anno[0])
                v = int(anno[1])
                if u < 0 or v < 0:
                    continue
                draw.ellipse((u - self.anno_radius, v - self.anno_radius, u + self.anno_radius, v + self.anno_radius),
                             fill=self.anno_color[i])
            del draw
            imgs[1] = np.asarray(img_21, np.uint8)
        return imgs