implement dataloader @2.2.0

autoaim/__init__.py

@@ -1,5 +1,6 @@
 from . import serial as aaserial
 from . import helpers
 from . import feature
+from . import dataloader
 from .feature import Feature
-from .feature import pipe
+from .feature import pipe

autoaim/dataloader.py

@@ -0,0 +1,190 @@
+# -*- coding: utf-8 -*-
+"""Data Loader Module
+
+This module save the features to a csv file.
+
+Author:
+    tccoin
+"""
+
+import os
+import csv
+import xml.etree.ElementTree as ET
+import cv2
+import numpy as np
+from toolz import pipe, curry
+from autoaim import *
+
+data_path = os.path.abspath(__file__ + '/../../data')
+
+
+class DataLoader():
+    '''@todo: finish implementation of __calcdict'''
+    __calcdict = {
+        'contour': {
+            'contour_len': lambda x: len(x)
+        },
+        'bounding_rect': {
+            'bounding_rect_w': lambda x: x[0],
+            'bounding_rect_y': lambda x: x[1],
+            'bounding_rect_ratio': lambda x: x[0]/x[1] if x[1] else 0
+        },
+        'bingo': {
+            'bingo': lambda x: int(x),
+        },
+    }
+
+    def __init__(self, datasets, feature_bucket, csv_filename, debug=False):
+        '''Example input:
+        datasets = ['test0', 'test1', ...]
+        feature_bucket = ['contour', 'bounding_rect', ...]
+        '''
+        self.feature_bucket = feature_bucket
+        self.csv_filename = csv_filename
+        self.debug = debug
+        for dataset in datasets:
+            self.load_dataset(dataset)
+
+    def load_dataset(self, dataset):
+        dataset_path = data_path+'/'+dataset
+        files = os.listdir(dataset_path)
+        self.new()
+        for file in files:
+            file_path = dataset_path+'/'+file
+            if os.path.isfile(file_path):
+                if self.load_img(dataset, file):
+                    break
+
+    def load_img(self, dataset, file):
+        '''return `exit`'''
+        # load labels
+        labeled_lamps, labeled_pairs = self.load_label(dataset, file)
+        # load features
+        img_path = os.path.join(data_path, dataset, file)
+        img = helpers.load(img_path)
+        # close preprocess to get more lamps
+        feature = Feature(img, preprocess=False)
+        lamps = feature.lamps
+        for feature_name in self.feature_bucket:
+            getattr(feature, feature_name+'s', '')
+        # label the real-time feature
+        for lamp in lamps:
+            lamp.bingo = False
+            for labeled_lamp in labeled_lamps:
+                if self.is_in(lamp.bounding_rect, labeled_lamp):
+                    lamp.bingo = True
+                    break
+        self.save(feature)
+        if self.debug:
+            exit = pipe(img.copy(),
+                        # feature.draw_contours,
+                        feature.draw_bounding_rects,
+                        self.draw_labeled_lamps()(labeled_lamps),
+                        self.draw_bingo_lamps()(feature),
+                        helpers.showoff
+                        )
+            return exit
+        else:
+            return False
+
+    def is_in(self, rect, labeled_rect):
+        margin = 15
+        x, y, w, h = rect
+        diff = abs(np.array([x, x+w, y, y+h]) - np.array(labeled_rect))
+        if len(np.where(diff > margin)[0]) == 0:
+            return True
+        else:
+            return False
+
+    def load_label(self, dataset, file):
+        label = os.path.splitext(file)[0]
+        label_path = os.path.join(
+            data_path, dataset, 'label', label+'.xml')
+        try:
+            tree = ET.ElementTree(file=label_path)
+        except:
+            raise Exception('Label file "{}"  not found!'.format(label_path))
+        root = tree.getroot()
+        lamps = []
+        pairs = []
+        for child in root:
+            if child.tag == 'object':
+                name = child[0].text
+                rect = (child[4][0], child[4][2], child[4][1],
+                        child[4][3])  # xmin,xmax,ymin,ymax
+                rect = [int(x.text) for x in rect]
+                if name == 'lamp':
+                    lamps.append(rect)
+                elif name == 'pair':
+                    pairs.append(rect)
+        return lamps, pairs
+
+    def new(self):
+        csv_filename = self.csv_filename
+        with open(data_path + '/'+csv_filename, 'w', newline='') as csvfile:
+            writer = csv.writer(csvfile, delimiter=' ',
+                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
+            writer.writerow(self.calc_row())
+
+    def save(self, feature):
+        csv_filename = self.csv_filename
+        with open(data_path + '/'+csv_filename, 'a', newline='') as csvfile:
+            writer = csv.writer(csvfile, delimiter=' ',
+                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
+            for lamp in feature.lamps:
+                writer.writerow(self.calc_row(lamp))
+
+    def calc_row(self, lamp=None):
+        calcdict = self.__calcdict
+        row = []
+        if lamp is None:
+            for feature_name in self.feature_bucket:
+                if feature_name in calcdict:
+                    row += calcdict[feature_name].keys()
+        else:
+            for feature_name in self.feature_bucket:
+                if feature_name in calcdict:
+                    for func in calcdict[feature_name].values():
+                        row.append(func(getattr(lamp, feature_name)))
+        return row
+
+    def draw_bingo_lamps(self):
+        '''Usage:dataloader.draw_bingo_lamps()(feature)'''
+        def draw(feature, img):
+            lamps = feature.lamps
+            boom_rects = [x.bounding_rect for x in lamps if not x.bingo]
+            bingo_rects = [x.bounding_rect for x in lamps if x.bingo]
+            for rect in boom_rects:
+                x, y, w, h = rect
+                cv2.rectangle(img, (x, y), (x+w, y+h), (200, 0, 200), 1)
+            for rect in bingo_rects:
+                x, y, w, h = rect
+                cv2.rectangle(img, (x, y), (x+w, y+h), (200, 0, 0), 2)
+            return img
+        return curry(draw)
+
+    def draw_labeled_lamps(self):
+        '''Usage:dataloader.draw_labeled_lamps()(rects)'''
+        def draw(rects, img):
+            for rect in rects:
+                xmin, xmax, ymin, ymax = rect
+                cv2.rectangle(img, (xmin, ymin), (xmax, ymax),
+                              (0, 0, 200), 3)
+            return img
+        return curry(draw)
+
+
+if __name__ == '__main__':
+    datasets = [
+        'test0',
+        # 'test1',
+    ]
+    feature_bucket = [
+        'contour',
+        'bounding_rect',
+        'rotated_rect',
+        'greyscale',
+        'point_area',
+        'bingo'
+    ]
+    dataloader = DataLoader(datasets, feature_bucket, 'test.csv', debug=True)

data/test.csv

@@ -0,0 +1,121 @@
+contour_len bounding_rect_w bounding_rect_y bounding_rect_ratio bingo
+45 176 455 0.3868131868131868 0
+66 400 443 0.9029345372460497 0
+19 406 451 0.9002217294900222 0
+13 331 251 1.3187250996015936 0
+19 269 244 1.1024590163934427 1
+20 197 240 0.8208333333333333 1
+21 450 203 2.2167487684729066 0
+14 243 193 1.2590673575129534 0
+10 428 172 2.488372093023256 0
+17 444 171 2.5964912280701755 0
+71 558 120 4.65 0
+24 576 103 5.592233009708738 0
+55 427 78 5.4743589743589745 0
+50 434 55 7.890909090909091 0
+106 193 48 4.020833333333333 0
+55 215 18 11.944444444444445 0
+10 323 405 0.7975308641975308 0
+15 285 393 0.7251908396946565 0
+16 370 390 0.9487179487179487 0
+32 490 355 1.380281690140845 0
+22 522 343 1.521865889212828 0
+8 451 331 1.3625377643504533 0
+44 498 329 1.513677811550152 0
+18 491 325 1.5107692307692309 0
+12 443 308 1.4383116883116882 0
+14 266 221 1.2036199095022624 1
+16 445 218 2.041284403669725 0
+24 384 217 1.7695852534562213 1
+24 202 216 0.9351851851851852 0
+7 212 104 2.0384615384615383 0
+10 435 98 4.438775510204081 0
+28 239 95 2.5157894736842104 0
+46 565 85 6.647058823529412 0
+58 240 84 2.857142857142857 0
+17 543 82 6.621951219512195 0
+9 284 82 3.4634146341463414 0
+13 572 71 8.056338028169014 0
+17 609 62 9.82258064516129 0
+10 265 57 4.649122807017544 0
+21 395 42 9.404761904761905 0
+24 264 35 7.542857142857143 0
+15 276 23 12.0 0
+17 368 16 23.0 0
+14 527 11 47.90909090909091 0
+18 368 8 46.0 0
+269 576 0 0 0
+66 607 1 607.0 0
+51 478 0 0 0
+94 416 445 0.9348314606741573 0
+7 562 423 1.3286052009456264 0
+50 213 340 0.6264705882352941 0
+15 312 251 1.2430278884462151 1
+12 308 249 1.2369477911646587 1
+23 240 247 0.97165991902834 1
+9 150 191 0.7853403141361257 0
+12 239 183 1.3060109289617485 0
+13 131 182 0.7197802197802198 0
+13 426 175 2.434285714285714 0
+24 445 174 2.557471264367816 0
+9 196 131 1.4961832061068703 0
+67 558 122 4.573770491803279 0
+20 578 105 5.504761904761905 0
+52 431 78 5.5256410256410255 0
+65 438 55 7.963636363636364 0
+9 216 52 4.153846153846154 0
+60 216 15 14.4 0
+31 377 246 1.532520325203252 1
+27 451 237 1.9029535864978904 1
+15 506 235 2.1531914893617023 0
+9 419 189 2.2169312169312168 0
+58 510 172 2.9651162790697674 0
+7 530 145 3.6551724137931036 0
+20 627 72 8.708333333333334 0
+98 242 60 4.033333333333333 0
+22 0 57 0.0 0
+63 481 52 9.25 0
+15 633 51 12.411764705882353 0
+31 15 38 0.39473684210526316 0
+51 251 28 8.964285714285714 0
+45 481 26 18.5 0
+12 18 0 0 0
+6 176 281 0.6263345195729537 1
+12 210 278 0.7553956834532374 1
+9 382 240 1.5916666666666666 1
+11 328 239 1.3723849372384938 1
+20 405 238 1.7016806722689075 0
+18 284 237 1.1983122362869199 0
+11 578 213 2.7136150234741785 0
+7 306 210 1.457142857142857 0
+10 130 188 0.6914893617021277 0
+14 262 186 1.4086021505376345 0
+8 456 177 2.5762711864406778 0
+11 448 176 2.5454545454545454 0
+50 567 123 4.609756097560975 0
+19 584 106 5.509433962264151 0
+53 435 82 5.304878048780488 0
+50 443 60 7.383333333333334 0
+116 196 55 3.5636363636363635 0
+43 217 23 9.434782608695652 0
+20 526 471 1.1167728237791932 0
+31 295 444 0.6644144144144144 0
+10 413 335 1.2328358208955223 1
+11 455 331 1.3746223564954683 1
+14 421 328 1.2835365853658536 1
+13 471 323 1.458204334365325 0
+20 391 316 1.2373417721518987 0
+14 295 294 1.0034013605442176 1
+14 270 292 0.9246575342465754 1
+12 411 290 1.4172413793103449 0
+13 314 289 1.0865051903114187 0
+23 254 282 0.900709219858156 0
+9 174 279 0.6236559139784946 1
+7 153 278 0.5503597122302158 1
+12 284 267 1.0636704119850187 0
+8 571 162 3.5246913580246915 0
+11 562 161 3.4906832298136647 0
+12 373 158 2.3607594936708862 0
+73 557 51 10.92156862745098 0
+29 585 34 17.205882352941178 0
+30 565 26 21.73076923076923 0

setup.py

@@ -9,7 +9,7 @@
 
 setup(
     name='AutoAim',
-    version='2.1.0',
+    version='2.2.0',
     description='A project for detecting armors.',
     author='FuXing PS',
     author_email='[email protected]',