From 9ff059935de72ac338374358f7430f0c0131d949 Mon Sep 17 00:00:00 2001 From: tccoin <1774241372@qq.com> Date: Tue, 26 Feb 2019 01:10:07 +0800 Subject: [PATCH] implement dataloader @2.2.0 --- autoaim/__init__.py | 3 +- autoaim/dataloader.py | 190 ++++++++++++++++++++++++++++++++++++++++++ data/test.csv | 121 +++++++++++++++++++++++++++ setup.py | 2 +- 4 files changed, 314 insertions(+), 2 deletions(-) create mode 100644 autoaim/dataloader.py create mode 100644 data/test.csv diff --git a/autoaim/__init__.py b/autoaim/__init__.py index d2938c6..6d7dda4 100644 --- a/autoaim/__init__.py +++ b/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 \ No newline at end of file +from .feature import pipe diff --git a/autoaim/dataloader.py b/autoaim/dataloader.py new file mode 100644 index 0000000..2967d0c --- /dev/null +++ b/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) diff --git a/data/test.csv b/data/test.csv new file mode 100644 index 0000000..3f125eb --- /dev/null +++ b/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 diff --git a/setup.py b/setup.py index a93f110..3e43a54 100755 --- a/setup.py +++ b/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='robovigor@gmail.com',