├── README.md ├── annotation_generator.py ├── annotation_gui.py ├── config.yml ├── config_dialog.py ├── create_train_data.py ├── crop.py ├── draw_bbox.py ├── image_dataset.py ├── image_util.py ├── inria_person_dataset.py ├── my_util.py └── train.py /README.md: -------------------------------------------------------------------------------- 1 | # HumanDetection 2 | Human detection program for Inria Person Dataset 3 | 4 | # Installation 5 | 6 | Download by git 7 | 8 | ``` 9 | git clone https://github.com/rupy/HumanDetection.git 10 | ``` 11 | 12 | or you can download as zip file from https://github.com/rupy/HumanDetection/archive/master.zip. 13 | 14 | ## Dependency 15 | 16 | You need environment like: 17 | 18 | ``` 19 | Python 2.7 20 | OpenCV 2.4.9 21 | Numpy 1.9.1 22 | PyYAML 3.11 23 | PySide 1.2.2 24 | ``` 25 | If you don't have Numpy and PyYAML, you can install these by pip like: 26 | 27 | ```Shell 28 | $ pip install numpy 29 | $ pip install pyyaml 30 | $ pip install pyside 31 | ``` 32 | 33 | ## Configure file 34 | 35 | Edit configure file, config.yml for your own file system. output directories are automatically created by program. Inria Person Dataset is in http://pascal.inrialpes.fr/data/human/. 36 | 37 | # Usage 38 | 39 | ## AnnotationGenerator 40 | 41 | ### Overview and configuration 42 | 43 | Annotation Generator is a GUI tool to generate annotation information for each image to learn dataset by opencv_traincascade. 44 | You don't need to use this class if you use Inria Person Dataset because it has annotation information in it. This program is for the dataset which has no annotation informations. 45 | 46 | You can create region information of objects to be detected. To use this class, you have to change config.yml for your own file system as follow: 47 | - pos_img_dir: Directory contains images to add annotations. 48 | - my_annotation_dir: Directory to save annotation infomations as pickle file. This directory is automatically created. 49 | - my_annotation_img_dir: Directory to save sample annotated images. This directory is automatically created. 50 | 51 | ### Simple code 52 | 53 | Write code as follow: 54 | 55 | ```python 56 | from annotation_generator import AnnotationGenerator 57 | 58 | # log level setting 59 | logging.root.setLevel(level=logging.INFO) 60 | 61 | # initialize AnnotationGenerator 62 | generator = AnnotationGenerator() 63 | 64 | # Do annotation work by GUI 65 | # if given True, generator skips file you already added annotations(default). 66 | # if given False, you can edit file you already added annotations. 67 | generator.generate_annotations(True) 68 | 69 | # create positive.dat for opencv 70 | generator.create_positive_dat() 71 | ``` 72 | 73 | or just run python like this: 74 | 75 | ```Shell 76 | $ python annotation_generator.py 77 | ``` 78 | ### How to annotate 79 | 80 | Drag point to point in image to create regions. You can use keys as follow: 81 | - [d] key: Delete a region added last. 82 | - [space] key: Save regions as pickle & go to next image. 83 | - [q] key: Quit annotation work. 84 | 85 | You can quit annotation work anytime. You can start from the image you quit brefore if you give True to generate_annotations(). 86 | 87 | ## AnnotationGUI 88 | 89 | ### Overview and configuration 90 | 91 | Annotation GUI is a GUI tool to generate annotation information for each image to learn dataset by opencv_traincascade. 92 | You don't need to use this class if you use Inria Person Dataset because it has annotation information in it. This program is for the dataset which has no annotation informations. 93 | 94 | ![annotation_tool](http://rupy.github.io/images/annotation_tool.jpg) 95 | 96 | You can create region information of objects to be detected. To use this class, you have to change config.yml for your own file system as follow: 97 | - pos_img_dir: Directory contains images to add annotations. 98 | - my_annotation_dir: Directory to save annotation infomations as pickle file. This directory is automatically created. 99 | - my_annotation_img_dir: Directory to save sample annotated images. This directory is automatically created. 100 | 101 | ### Simple code 102 | 103 | Write code as follow: 104 | 105 | ```python 106 | from annotation_gui import AnnotationGUI 107 | 108 | # log level setting 109 | logging.root.setLevel(level=logging.INFO) 110 | 111 | # run app 112 | app = QtGui.QApplication(sys.argv) 113 | ag = AnnotationGUI() 114 | ag.show() 115 | sys.exit(app.exec_()) 116 | ``` 117 | or just run python like this: 118 | 119 | ```Shell 120 | $ python annotation_gui.py 121 | ``` 122 | ### How to annotate 123 | 124 | First, choose file name from list, then image will be loaded. Drag point to point in image to create regions. You can save and undo by button. Green icon in the list shows file has annotation information already and red shows no information. 125 | 126 | ## InriaPersonDataSet 127 | 128 | ## ImageDataSet 129 | -------------------------------------------------------------------------------- /annotation_generator.py: -------------------------------------------------------------------------------- 1 | #coding: utf-8 2 | 3 | import os 4 | import cv2 5 | import logging 6 | import sys 7 | import yaml 8 | try: 9 | import cPickle as pickle 10 | except: 11 | import pickle 12 | 13 | class AnnotationGenerator: 14 | 15 | CONFIG_YAML = 'config.yml' 16 | GENERATOR_WINDOW_NAME = 'generator' 17 | 18 | def __init__(self): 19 | 20 | # log setting 21 | program = os.path.basename(sys.argv[0]) 22 | self.logger = logging.getLogger(program) 23 | logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s') 24 | 25 | # load config file 26 | f = open(self.CONFIG_YAML, 'r') 27 | self.config = yaml.load(f) 28 | f.close() 29 | 30 | # set dataset path 31 | self.pos_img_dir = self.config['dataset']['pos_img_dir'] 32 | 33 | # set output path 34 | self.my_annotation_dir = self.config['output']['my_annotation_dir'] 35 | self.my_annotation_img_dir = self.config['output']['my_annotation_img_dir'] 36 | 37 | # create output paths 38 | if not os.path.isdir(self.my_annotation_dir): 39 | os.makedirs(self.my_annotation_dir) 40 | if not os.path.isdir(self.my_annotation_img_dir): 41 | os.makedirs(self.my_annotation_img_dir) 42 | 43 | # set array of all file names 44 | self.my_annotation_files = [file_name for file_name in os.listdir(self.my_annotation_dir) if not file_name.startswith('.')] 45 | self.my_annotation_files.sort() 46 | self.pos_img_files = [file_name for file_name in os.listdir(self.pos_img_dir) if not file_name.startswith('.')] 47 | self.pos_img_files.sort() 48 | 49 | # initialize mouse event 50 | cv2.namedWindow(self.GENERATOR_WINDOW_NAME) 51 | cv2.setMouseCallback(self.GENERATOR_WINDOW_NAME, self.on_mouse) 52 | 53 | # mouse location 54 | self.im_orig = None 55 | self.start_pt = (0, 0) 56 | self.end_pt = (0, 0) 57 | self.mouse_dragging = False 58 | self.bboxes = [] 59 | 60 | def on_mouse(self, event, x, y, flags, param): 61 | 62 | x = min(max(x, 0), self.im_orig.shape[1] - 1) 63 | y = min(max(y, 0), self.im_orig.shape[0] - 1) 64 | 65 | if event == cv2.EVENT_LBUTTONDOWN: 66 | self.logger.info('DOWN: %d, %d', x, y) 67 | self.start_pt = (x, y) 68 | self.end_pt = (x, y) 69 | self.mouse_dragging = True 70 | elif event == cv2.EVENT_LBUTTONUP: 71 | self.logger.info('UP: %d, %d', x, y) 72 | self.end_pt = (x, y) 73 | self.bboxes.append((self.start_pt, self.end_pt)) 74 | self.start_pt = self.end_pt = (0, 0) 75 | self.mouse_dragging = False 76 | elif event == cv2.EVENT_MOUSEMOVE and self.mouse_dragging: 77 | # self.logger.info('DRAG: %d, %d', x, y) 78 | self.end_pt = (x, y) 79 | 80 | 81 | def generate_my_annotation(self, img_path, edit=False): 82 | 83 | # annotation path 84 | head, tail = os.path.split(img_path) 85 | # root, ext = os.path.splitext(tail) 86 | annotation_path = self.my_annotation_dir + tail + '.pkl' 87 | 88 | # bbox path 89 | bbox_path = self.my_annotation_img_dir + 'bbox_' + tail 90 | 91 | # load image 92 | self.im_orig = cv2.imread(img_path) 93 | 94 | # if edit is true, load bbox info from annotation file 95 | if edit: 96 | f = open(annotation_path, 'rb') 97 | self.bboxes = pickle.load(f) 98 | f.close() 99 | 100 | while True: 101 | im_copy = self.im_orig.copy() 102 | 103 | # draw rectangles 104 | if self.start_pt is not (0, 0) and self.end_pt is not (0, 0): 105 | cv2.rectangle(im_copy, self.start_pt, self.end_pt, (0, 0, 255), 1) 106 | for box in self.bboxes: 107 | cv2.rectangle(im_copy, box[0], box[1], (0, 255, 0), 1) 108 | 109 | # show image to generate annotations 110 | cv2.imshow(self.GENERATOR_WINDOW_NAME, im_copy) 111 | key = cv2.waitKey(10) 112 | if key == ord('q'): # 'q' key 113 | cv2.destroyAllWindows() 114 | return False 115 | elif key == 32: # space key 116 | self.logger.info('saving annotation data: %s', annotation_path) 117 | f = open(annotation_path, 'wb') 118 | pickle.dump(self.bboxes, f) 119 | f.close() 120 | self.logger.info('saving bounding box data: %s', bbox_path) 121 | cv2.imwrite(bbox_path, im_copy) 122 | self.bboxes = [] 123 | return True 124 | elif key == ord('d'): # 'd' key 125 | if len(self.bboxes) > 0: 126 | self.bboxes.pop() 127 | else: 128 | self.logger.info('no bounding boxes to delete') 129 | return True 130 | 131 | def generate_annotations(self, skip=True): 132 | 133 | for pos_image_file in self.pos_img_files: 134 | 135 | edit = False 136 | if pos_image_file in [os.path.splitext(annotation_file)[0] for annotation_file in self.my_annotation_files]: 137 | if skip: 138 | self.logger.info('skipping: %s is already added annotation', pos_image_file) 139 | continue 140 | else: 141 | self.logger.info('edit: %s is already added annotation', pos_image_file) 142 | edit = True 143 | else: 144 | self.logger.info('new: %s', pos_image_file) 145 | 146 | pos_img_path = self.pos_img_dir + pos_image_file 147 | is_continue = self.generate_my_annotation(pos_img_path, edit) 148 | if not is_continue: 149 | return 150 | 151 | def create_positive_dat(self): 152 | output_text = "" 153 | self.logger.info("begin creating positive.dat") 154 | for file_name in self.my_annotation_files: 155 | 156 | # annotation path 157 | annotation_path = self.my_annotation_dir + file_name 158 | f = open(annotation_path, 'rb') 159 | bboxes = pickle.load(f) 160 | f.close() 161 | root, ext = os.path.splitext(file_name) 162 | output_text += "%s %d " % (self.pos_img_dir + root, len(bboxes)) 163 | for bbox in bboxes: 164 | x_min, y_min = min(bbox[0][0], bbox[1][0]), min(bbox[0][1], bbox[1][1]) 165 | x_max, y_max = max(bbox[0][0], bbox[1][0]), max(bbox[0][1], bbox[1][1]) 166 | w = x_max - x_min 167 | h = y_max - y_min 168 | output_text += "%d %d %d %d " % (x_min, y_min, w, h) 169 | output_text += "\n" 170 | # print output_text 171 | self.logger.info("writing data to positive.dat") 172 | f = open('positive.dat', 'w') 173 | f.write(output_text) 174 | f.close() 175 | self.logger.info("completed writing data to positive.dat") 176 | 177 | if __name__ == '__main__': 178 | 179 | # log level setting 180 | logging.root.setLevel(level=logging.INFO) 181 | 182 | # generate AnnotationGenerator 183 | generator = AnnotationGenerator() 184 | 185 | # generate annotations by GUI 186 | # if given True, generator skips file you already added annotations(default). 187 | # if given False, you can edit file you already added annotations. 188 | generator.generate_annotations(True) 189 | 190 | # create positive.dat for opencv 191 | generator.create_positive_dat() 192 | -------------------------------------------------------------------------------- /annotation_gui.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | import os 5 | import cv2 6 | import logging 7 | from PySide import QtCore, QtGui 8 | try: 9 | import cPickle as pickle 10 | except: 11 | import pickle 12 | import image_dataset 13 | import config_dialog 14 | 15 | class AnnotationGUI(QtGui.QMainWindow): 16 | 17 | def __init__(self): 18 | 19 | # log setting 20 | program = os.path.basename(sys.argv[0]) 21 | self.logger = logging.getLogger(program) 22 | logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s') 23 | 24 | super(AnnotationGUI, self).__init__() 25 | self.setWindowTitle('Annotation Tool') 26 | 27 | self.dataset = image_dataset.ImageDataSet() 28 | 29 | 30 | self.dialog = None 31 | if self.dataset.pos_img_dir == ''\ 32 | or self.dataset.output_dir == ''\ 33 | or not os.path.isdir(self.dataset.pos_img_dir)\ 34 | or not os.path.isdir(self.dataset.output_dir): 35 | self.open_dir_dialog() 36 | 37 | self.cv_img = None 38 | self.cv_bbox_img = None 39 | self.start_pt = None 40 | self.end_pt = None 41 | self.bboxes = [] 42 | 43 | self.resize(800, 600) 44 | self.__init_menu_bar() 45 | self.__init_tool_bar() 46 | self.__init_status_bar() 47 | self.__init_ui() 48 | 49 | 50 | 51 | def __init_ui(self): 52 | 53 | self.image_label = QtGui.QLabel(self) 54 | self.image_label.setBackgroundRole(QtGui.QPalette.Base) 55 | self.image_label.installEventFilter(self) 56 | 57 | self.list_widget = QtGui.QListWidget(self) 58 | self.list_widget.itemSelectionChanged.connect(self.open) 59 | 60 | self.list_widget.setGeometry(10, 50, 200, 500) 61 | self.image_label.move(220, 50) 62 | 63 | self.__init_list_widget() 64 | 65 | def __init_menu_bar(self): 66 | menu_bar = QtGui.QMenuBar() 67 | file_menu = QtGui.QMenu('File', self) 68 | open_action = file_menu.addAction('Open') 69 | open_action.triggered.connect(self.open_dir_dialog) 70 | exit_action = file_menu.addAction('Close Annotation Tool') 71 | exit_action.setShortcut('Ctrl+Q') 72 | exit_action.triggered.connect(QtGui.qApp.quit) 73 | menu_bar.addMenu(file_menu) 74 | self.setMenuBar(menu_bar) 75 | 76 | def __init_tool_bar(self): 77 | open_action = QtGui.QAction(self.style().standardIcon(QtGui.QStyle.SP_DirOpenIcon), 'Open', self) 78 | open_action.triggered.connect(self.open_dir_dialog) 79 | open_action.setShortcut('Ctrl+O') 80 | 81 | save_action = QtGui.QAction(self.style().standardIcon(QtGui.QStyle.SP_DialogSaveButton), 'Save', self) 82 | save_action.triggered.connect(self.save) 83 | save_action.setShortcut('Ctrl+S') 84 | 85 | undo_action = QtGui.QAction(self.style().standardIcon(QtGui.QStyle.SP_TrashIcon), 'Undo', self) 86 | undo_action.triggered.connect(self.remove_box) 87 | undo_action.setShortcut('Ctrl+Z') 88 | 89 | count_action = QtGui.QAction(self.style().standardIcon(QtGui.QStyle.SP_MessageBoxInformation), 'Count', self) 90 | count_action.triggered.connect(self.count_all_bbox) 91 | count_action.setShortcut('Ctrl+C') 92 | 93 | self.toolbar = self.addToolBar('Toolbar') 94 | self.toolbar.addAction(open_action) 95 | self.toolbar.addAction(save_action) 96 | self.toolbar.addAction(undo_action) 97 | self.toolbar.addAction(count_action) 98 | 99 | def __init_status_bar(self): 100 | self.status_label = QtGui.QLabel() 101 | self.status_label.setText('') 102 | self.status_bar = self.statusBar() 103 | self.status_bar.addWidget(self.status_label) 104 | 105 | def __init_list_widget(self): 106 | 107 | # add list items 108 | self.list_widget.clear() 109 | self.list_widget.addItems(self.dataset.pos_img_files) 110 | 111 | # set icon 112 | for i, is_made in enumerate(self.dataset.get_annotation_existence_list()): 113 | if is_made: 114 | self.list_widget.item(i).setIcon(self.style().standardIcon(QtGui.QStyle.SP_DialogYesButton)) 115 | else: 116 | self.list_widget.item(i).setIcon(self.style().standardIcon(QtGui.QStyle.SP_DialogNoButton)) 117 | self.list_widget.setCurrentRow(0) 118 | 119 | def open_dir_dialog(self): 120 | self.dialog = config_dialog.ConfigDialig(self, self.dataset) 121 | result = self.dialog.exec_() 122 | if result == QtGui.QDialog.Accepted: 123 | self.dataset.save_config() 124 | self.__init_list_widget() 125 | self.open() 126 | elif result == QtGui.QDialog.Rejected: 127 | pass 128 | print result 129 | 130 | def count_all_bbox(self): 131 | 132 | # show count in message box 133 | msgBox = QtGui.QMessageBox() 134 | msgBox.setText("all boxes count: %d" % self.dataset.count_all_bboxes()) 135 | msgBox.exec_() 136 | 137 | def convert_cv_img2qt_img(self, cv_img): 138 | height, width, dim = cv_img.shape 139 | bytes_per_line = dim * width 140 | qt_img = QtGui.QImage(cv_img.data, width, height, bytes_per_line, QtGui.QImage.Format_RGB888) 141 | qt_img_rgb = qt_img.rgbSwapped() # BGR to RGB 142 | return qt_img_rgb 143 | 144 | def load_image(self, img_file): 145 | 146 | # read image 147 | self.cv_img = self.dataset.read_img(img_file) 148 | 149 | # draw bounding boxes 150 | self.draw_bbox() 151 | self.set_image_label() 152 | 153 | # set labels 154 | self.status_label.setText("box num: %d" % len(self.bboxes)) 155 | 156 | def eventFilter(self, source, event): 157 | # drag start 158 | if event.type() == QtCore.QEvent.MouseButtonPress and source is self.image_label: 159 | if event.button() == QtCore.Qt.LeftButton: 160 | pos = event.pos() 161 | x = min(max(pos.x(), 0), self.cv_img.shape[1] - 1) 162 | y = min(max(pos.y(), 0), self.cv_img.shape[0] - 1) 163 | pt = (x, y) 164 | self.start_pt = pt 165 | # print "Drag start (%d, %d)" % pt 166 | # dragging 167 | if event.type() == QtCore.QEvent.MouseMove and source is self.image_label: 168 | if event.buttons() & QtCore.Qt.LeftButton: # use buttons() instead of button() 169 | pos = event.pos() 170 | x = min(max(pos.x(), 0), self.cv_img.shape[1] - 1) 171 | y = min(max(pos.y(), 0), self.cv_img.shape[0] - 1) 172 | pt = (x, y) 173 | self.end_pt = pt 174 | # print "Dragging (%d, %d)" % pt 175 | self.draw_bbox() 176 | self.set_image_label() 177 | # drag end 178 | if event.type() == QtCore.QEvent.MouseButtonRelease and source is self.image_label: 179 | if event.button() == QtCore.Qt.LeftButton: 180 | if self.start_pt is not None: 181 | pos = event.pos() 182 | x = min(max(pos.x(), 0), self.cv_img.shape[1] - 1) 183 | y = min(max(pos.y(), 0), self.cv_img.shape[0] - 1) 184 | pt = (x, y) 185 | self.end_pt = pt 186 | self.bboxes.append((self.start_pt, self.end_pt)) 187 | self.start_pt = self.end_pt = None 188 | # print "Drag end (%d, %d)" % pt 189 | self.draw_bbox() 190 | self.set_image_label() 191 | self.status_label.setText("box num: %d" % len(self.bboxes)) 192 | 193 | return QtGui.QWidget.eventFilter(self, source, event) 194 | 195 | def draw_bbox(self): 196 | if self.cv_img is not None: 197 | self.cv_bbox_img = self.dataset.draw_areas(self.cv_img, self.start_pt, self.end_pt, self.bboxes) 198 | 199 | def set_image_label(self): 200 | if self.cv_bbox_img is not None: 201 | qt_img = self.convert_cv_img2qt_img(self.cv_bbox_img) 202 | self.image_label.setPixmap(QtGui.QPixmap.fromImage(qt_img)) 203 | self.image_label.adjustSize() 204 | 205 | def save(self): 206 | 207 | # image path 208 | idx = self.list_widget.currentRow() 209 | img_file = self.dataset.pos_img_files[idx] 210 | 211 | # save 212 | self.dataset.save_annotation(img_file, self.bboxes) 213 | self.dataset.write_bbox_img(img_file, self.cv_bbox_img) 214 | 215 | # reload annotation files 216 | self.dataset.reload_my_annotation_files() 217 | 218 | # change current icon yes 219 | self.list_widget.item(idx).setIcon(self.style().standardIcon(QtGui.QStyle.SP_DialogYesButton)) 220 | 221 | # show msg box 222 | msgBox = QtGui.QMessageBox() 223 | msgBox.setText("saved") 224 | msgBox.exec_() 225 | 226 | def remove_box(self): 227 | # get current image path 228 | idx = self.list_widget.currentRow() 229 | img_file = self.dataset.pos_img_files[idx] 230 | 231 | # delete last box 232 | if len(self.bboxes) > 0: 233 | self.bboxes.pop() 234 | print self.bboxes 235 | else: 236 | self.logger.info('no bounding boxes to delete') 237 | self.load_image(img_file) 238 | 239 | def open(self): 240 | # get current image path 241 | idx = self.list_widget.currentRow() 242 | img_file = self.dataset.pos_img_files[idx] 243 | 244 | # load 245 | self.bboxes = self.dataset.load_annotation(img_file) 246 | self.load_image(img_file) 247 | 248 | if __name__ == '__main__': 249 | 250 | import sys 251 | 252 | # log level setting 253 | logging.root.setLevel(level=logging.INFO) 254 | 255 | app = QtGui.QApplication(sys.argv) 256 | # app.setStyle(QtGui.QStyleFactory.create('Cleanlooks')) 257 | ag = AnnotationGUI() 258 | ag.show() 259 | sys.exit(app.exec_()) -------------------------------------------------------------------------------- /config.yml: -------------------------------------------------------------------------------- 1 | dataset: 2 | pos_img_dir: ./INRIAPerson/Train/pos/ 3 | neg_img_dir: ./INRIAPerson/Train/neg/ 4 | annotation_dir: ./INRIAPerson/Train/annotations/ 5 | test_img_dir: ./INRIAPerson/Test/pos/ 6 | output: 7 | output_dir: ./output/ 8 | # cropped_dir: ./output/cropped/ 9 | # bounding_box_out_dir: ./output/bounding_box/ 10 | # out_dir: ./output/out/ 11 | # cascade_xml_dir: ./output/hog3/ 12 | # my_annotation_dir: ./output/my_annotation/ 13 | # my_annotation_img_dir: ./output/my_annotation_img/ -------------------------------------------------------------------------------- /config_dialog.py: -------------------------------------------------------------------------------- 1 | from PySide import QtCore, QtGui 2 | 3 | class ConfigDialig(QtGui.QDialog): 4 | 5 | def __init__(self, parent, dataset): 6 | super(ConfigDialig, self).__init__(parent) 7 | 8 | self.dataset = dataset 9 | self.tmp_config = self.dataset.config 10 | 11 | self.setWindowTitle('title') 12 | # self.setWindowFlags(QtCore.Qt.Tool) 13 | 14 | self.layout = QtGui.QVBoxLayout() 15 | 16 | self.pos_label = QtGui.QLabel() 17 | self.pos_label.setText('Positive Directory') 18 | self.layout.addWidget(self.pos_label) 19 | self.pos_dir_line_edit = QtGui.QLineEdit(self.dataset.pos_img_dir) 20 | self.layout.addWidget(self.pos_dir_line_edit) 21 | self.pos_dir_button = QtGui.QPushButton('Browse') 22 | self.pos_dir_button.clicked.connect(self.set_config_pos) 23 | self.layout.addWidget(self.pos_dir_button) 24 | 25 | self.neg_label = QtGui.QLabel() 26 | self.neg_label.setText('Negative Directory') 27 | self.layout.addWidget(self.neg_label) 28 | self.neg_dir_line_edit = QtGui.QLineEdit(self.dataset.neg_img_dir) 29 | self.layout.addWidget(self.neg_dir_line_edit) 30 | self.neg_dir_button = QtGui.QPushButton('Browse') 31 | self.neg_dir_button.clicked.connect(self.set_config_neg) 32 | self.layout.addWidget(self.neg_dir_button) 33 | 34 | self.test_label = QtGui.QLabel() 35 | self.test_label.setText('Test Directory') 36 | self.layout.addWidget(self.test_label) 37 | self.test_dir_line_edit = QtGui.QLineEdit(self.dataset.test_img_dir) 38 | self.layout.addWidget(self.test_dir_line_edit) 39 | self.test_dir_button = QtGui.QPushButton('Browse') 40 | self.test_dir_button.clicked.connect(self.set_config_test) 41 | self.layout.addWidget(self.test_dir_button) 42 | 43 | self.output_label = QtGui.QLabel() 44 | self.output_label.setText('Output Directory') 45 | self.layout.addWidget(self.output_label) 46 | self.output_dir_line_edit = QtGui.QLineEdit(self.dataset.output_dir) 47 | self.layout.addWidget(self.output_dir_line_edit) 48 | self.output_dir_button = QtGui.QPushButton('Browse') 49 | self.output_dir_button.clicked.connect(self.set_config_output) 50 | self.layout.addWidget(self.output_dir_button) 51 | 52 | self.button_box = QtGui.QDialogButtonBox(QtGui.QDialogButtonBox.Ok| QtGui.QDialogButtonBox.Cancel) 53 | # self.dialog_button.addButton(ok_button, QtGui.QDialogButtonBox.ActionRole) 54 | # self.dialog_button.addButton(ng_button, QtGui.QDialogButtonBox.ActionRole) 55 | self.button_box.accepted.connect(self.accept) 56 | self.button_box.rejected.connect(self.reject) 57 | self.layout.addWidget(self.button_box) 58 | self.setLayout(self.layout) 59 | 60 | def set_config_pos(self): 61 | get_dir = QtGui.QFileDialog.getExistingDirectory(self, "Open Directory") 62 | if get_dir: 63 | if not get_dir.endswith('/'): 64 | get_dir += '/' 65 | self.pos_dir_line_edit.setText(get_dir) 66 | self.tmp_config['dataset']['pos_img_dir'] = get_dir 67 | 68 | def set_config_neg(self): 69 | get_dir = QtGui.QFileDialog.getExistingDirectory(self, "Open Directory") 70 | if get_dir: 71 | if not get_dir.endswith('/'): 72 | get_dir += '/' 73 | self.neg_dir_line_edit.setText(get_dir) 74 | self.tmp_config['dataset']['neg_img_dir'] = get_dir 75 | 76 | def set_config_test(self): 77 | get_dir = QtGui.QFileDialog.getExistingDirectory(self, "Open Directory") 78 | if get_dir: 79 | if not get_dir.endswith('/'): 80 | get_dir += '/' 81 | self.test_dir_line_edit.setText(get_dir) 82 | self.tmp_config['dataset']['test_img_dir'] = get_dir 83 | 84 | def set_config_output(self): 85 | get_dir = QtGui.QFileDialog.getExistingDirectory(self, "Open Directory") 86 | if get_dir: 87 | if not get_dir.endswith('/'): 88 | get_dir += '/' 89 | self.output_dir_line_edit.setText(get_dir) 90 | self.tmp_config['dataset']['output_img_dir'] = get_dir 91 | -------------------------------------------------------------------------------- /create_train_data.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | from image_dataset import ImageDataSet 5 | import logging 6 | 7 | if __name__ == '__main__': 8 | 9 | logging.root.setLevel(level=logging.INFO) 10 | dataset = ImageDataSet() 11 | dataset.create_positive_dat_with_my_annotation() 12 | dataset.create_negative_dat() 13 | dataset.create_samples(use_my_annotation=True, width=24, height=24) 14 | 15 | -------------------------------------------------------------------------------- /crop.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | from image_dataset import ImageDataSet 5 | import logging 6 | 7 | if __name__ == '__main__': 8 | 9 | logging.root.setLevel(level=logging.INFO) 10 | dataset = ImageDataSet() 11 | dataset.create_crop_with_my_annotation() 12 | 13 | -------------------------------------------------------------------------------- /draw_bbox.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | from image_dataset import ImageDataSet 5 | import logging 6 | 7 | if __name__ == '__main__': 8 | 9 | logging.root.setLevel(level=logging.INFO) 10 | dataset = ImageDataSet() 11 | dataset.draw_bounding_boxes_for_all() 12 | -------------------------------------------------------------------------------- /image_dataset.py: -------------------------------------------------------------------------------- 1 | #coding: utf-8 2 | 3 | import os 4 | import cv2 5 | import image_util as iu 6 | import logging 7 | import sys 8 | import yaml 9 | import subprocess 10 | import my_util 11 | 12 | class ImageDataSet: 13 | 14 | CONFIG_YAML = 'nail_config.yml' 15 | 16 | OUT_DIR = 'outdir/' 17 | CROPPED_DIR = 'cropped/' 18 | CASCADE_XML_DIR = 'hog/' 19 | MY_ANNOTATION_DIR = 'my_annotation/' 20 | MY_ANNOTATION_IMG_DIR = 'bbox/' 21 | 22 | def __init__(self): 23 | 24 | # log setting 25 | program = os.path.basename(sys.argv[0]) 26 | self.logger = logging.getLogger(program) 27 | logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s') 28 | 29 | # load config file 30 | f = open(self.CONFIG_YAML, 'r') 31 | self.config = yaml.load(f) 32 | f.close() 33 | 34 | self.__init_dir() 35 | 36 | self.cascade = None 37 | 38 | def __init_dir(self): 39 | 40 | # assign empty value if dictionary is not set 41 | if 'dataset' not in self.config: 42 | self.config['dataset'] = None 43 | if 'pos_img_dir' not in self.config['dataset']: 44 | self.config['dataset']['pos_img_dir'] = '' 45 | if 'neg_img_dir' not in self.config['dataset']: 46 | self.config['dataset']['neg_img_dir'] = '' 47 | if 'test_img_dir' not in self.config['dataset']: 48 | self.config['dataset']['test_img_dir'] = '' 49 | if 'output' not in self.config: 50 | self.config['output'] = None 51 | if 'output_dir' not in self.config['output']: 52 | self.config['dataset']['output_dir'] = '' 53 | 54 | # set dataset path 55 | self.pos_img_dir = self.config['dataset']['pos_img_dir'] 56 | self.neg_img_dir = self.config['dataset']['neg_img_dir'] 57 | self.test_img_dir = self.config['dataset']['test_img_dir'] 58 | 59 | # set output path 60 | self.output_dir = self.config['output']['output_dir'] 61 | self.out_dir = self.output_dir + self.OUT_DIR 62 | self.cascade_xml_dir = self.output_dir + self.CASCADE_XML_DIR 63 | self.my_annotation_dir = self.output_dir + self.MY_ANNOTATION_DIR 64 | self.cropped_dir = self.output_dir + self.CROPPED_DIR 65 | self.my_annotation_img_dir = self.output_dir + self.MY_ANNOTATION_IMG_DIR 66 | 67 | # create output paths 68 | if not os.path.isdir(self.out_dir): 69 | os.makedirs(self.out_dir) 70 | if not os.path.isdir(self.cascade_xml_dir): 71 | os.makedirs(self.cascade_xml_dir) 72 | if not os.path.isdir(self.my_annotation_dir): 73 | os.makedirs(self.my_annotation_dir) 74 | if not os.path.isdir(self.cropped_dir): 75 | os.makedirs(self.cropped_dir) 76 | if not os.path.isdir(self.my_annotation_img_dir): 77 | os.makedirs(self.my_annotation_img_dir) 78 | 79 | # set array of all file names 80 | self.pos_img_files = [file_name for file_name in os.listdir(self.pos_img_dir) if not file_name.startswith('.')] 81 | self.pos_img_files.sort() 82 | self.neg_img_files = [file_name for file_name in os.listdir(self.neg_img_dir) if not file_name.startswith('.')] 83 | self.neg_img_files.sort() 84 | self.test_img_files = [file_name for file_name in os.listdir(self.test_img_dir) if not file_name.startswith('.')] 85 | self.test_img_files.sort() 86 | self.my_annotation_files = [file_name for file_name in os.listdir(self.my_annotation_dir) if not file_name.startswith('.')] 87 | self.my_annotation_files.sort() 88 | self.cropped_files = [file_name for file_name in os.listdir(self.cropped_dir) if not file_name.startswith('.')] 89 | self.cropped_files.sort() 90 | 91 | def save_config(self): 92 | 93 | # save config.yml 94 | self.logger.info("saving config file") 95 | f = open(self.CONFIG_YAML, 'w') 96 | f.write(yaml.dump(self.config, default_flow_style=False)) 97 | f.close() 98 | 99 | # reload valuables related directories 100 | self.__init_dir() 101 | 102 | def create_crop_with_my_annotation(self): 103 | 104 | self.logger.info("begin creating crop file") 105 | for annotation_file_name in self.my_annotation_files: 106 | 107 | # read image 108 | img_file_name = os.path.splitext(annotation_file_name)[0] 109 | img = cv2.imread(self.pos_img_dir + img_file_name) 110 | self.logger.info("cropping: %s", img_file_name) 111 | 112 | # read my annotation 113 | annotation_path = self.my_annotation_dir + annotation_file_name 114 | bboxes = my_util.my_unpickle(annotation_path) 115 | 116 | # crop 117 | for i, box in enumerate(bboxes): 118 | left_up_x = min(box[0][0], box[1][0]) 119 | left_up_y = min(box[0][1], box[1][1]) 120 | right_down_x = max(box[0][0], box[1][0]) 121 | right_down_y = max(box[0][1], box[1][1]) 122 | im_crop = iu.image_crop(img, (left_up_x, left_up_y), (right_down_x, right_down_y)) 123 | root, ext = os.path.splitext(img_file_name) 124 | crop_file_name = root + str(i) + ext 125 | cv2.imwrite(self.cropped_dir + crop_file_name, im_crop) 126 | self.logger.info("completed creating crop file") 127 | 128 | def create_positive_dat_with_my_annotation(self): 129 | output_text = "" 130 | self.logger.info("begin creating positive.dat") 131 | for file_name in self.my_annotation_files: 132 | 133 | # annotation path 134 | annotation_path = self.my_annotation_dir + file_name 135 | bboxes = my_util.my_unpickle(annotation_path) 136 | root, ext = os.path.splitext(file_name) 137 | output_text += "%s %d " % (self.pos_img_dir + root, len(bboxes)) 138 | for bbox in bboxes: 139 | x_min, y_min = min(bbox[0][0], bbox[1][0]), min(bbox[0][1], bbox[1][1]) 140 | x_max, y_max = max(bbox[0][0], bbox[1][0]), max(bbox[0][1], bbox[1][1]) 141 | w = x_max - x_min 142 | h = y_max - y_min 143 | output_text += "%d %d %d %d " % (x_min, y_min, w, h) 144 | output_text += "\n" 145 | 146 | self.logger.info("writing data to positive.dat") 147 | f = open('positive.dat', 'w') 148 | f.write(output_text) 149 | f.close() 150 | self.logger.info("completed writing data to positive.dat") 151 | 152 | def create_positive_dat_by_image_size(self): 153 | output_text = "" 154 | self.logger.info("begin creating positive.dat") 155 | for file_name in self.pos_img_files: 156 | 157 | file_path = self.pos_img_dir + file_name 158 | im = cv2.imread(file_path) 159 | output_text += "%s %d " % (file_path, 1) 160 | output_text += "%d %d %d %d \n" % (0, 0, im.shape[0], im.shape[1]) 161 | self.logger.info("writing data to positive.dat") 162 | f = open('positive.dat', 'w') 163 | f.write(output_text) 164 | f.close() 165 | self.logger.info("completed writing data to positive.dat") 166 | 167 | def create_samples(self, use_my_annotation=False, width=24, height=24): 168 | 169 | if use_my_annotation: 170 | self.create_positive_dat_with_my_annotation() 171 | else: 172 | self.create_positive_dat_by_image_size() 173 | self.create_negative_dat() 174 | 175 | params = { 176 | 'info': 'positive.dat', 177 | 'vec': 'positive.vec', 178 | 'num': len(self.pos_img_files), 179 | 'width': width, 180 | 'height': height 181 | } 182 | cmd = "opencv_createsamples -info %(info)s -vec %(vec)s -num %(num)d -w %(width)d -h %(height)d" % params 183 | self.logger.info("running command: %s", cmd) 184 | subprocess.call(cmd.strip().split(" ")) 185 | 186 | def create_negative_dat(self): 187 | output_text = "" 188 | self.logger.info("begin creating negative.dat") 189 | for file_name in self.neg_img_files: 190 | 191 | file_path = self.neg_img_dir + file_name 192 | output_text += file_path 193 | output_text += "\n" 194 | self.logger.info("writing data to negative.dat") 195 | f = open('negative.dat', 'w') 196 | f.write(output_text) 197 | f.close() 198 | self.logger.info("completed writing data to negative.dat") 199 | 200 | def train_cascade(self, feature_type='HOG', max_false_alarm_rate=0.4, min_hit_rate=0.995 ,width=24, height=24, pos_rate=0.8): 201 | pos_size = len(self.pos_img_files) 202 | neg_size = len(self.neg_img_files) 203 | 204 | params = { 205 | 'data': self.cascade_xml_dir, 206 | 'vec': 'positive.vec', 207 | 'bg': 'negative.dat', 208 | 'num_pos': pos_size * pos_rate, 209 | 'num_neg': neg_size, 210 | 'feature_type': feature_type, 211 | 'min_hit_rate': min_hit_rate, 212 | 'max_false_alarm_rate': max_false_alarm_rate, 213 | 'width': width, 214 | 'height': height 215 | } 216 | 217 | cmd = "opencv_traincascade -data %(data)s -vec %(vec)s -bg %(bg)s -numPos %(num_pos)d -numNeg %(num_neg)d -featureType %(feature_type)s -minHitRate %(min_hit_rate)s -maxFalseAlarmRate %(max_false_alarm_rate)s -w %(width)d -h %(height)d" % params 218 | 219 | self.logger.info("running command: %s", cmd) 220 | subprocess.call(cmd.strip().split(" ")) 221 | 222 | def inside(self, r, q): 223 | rx, ry, rw, rh = r 224 | qx, qy, qw, qh = q 225 | return rx > qx and ry > qy and rx + rw < qx + qw and ry + rh < qy + qh 226 | 227 | 228 | def draw_detections(self, img, rects, thickness = 1): 229 | for x, y, w, h in rects: 230 | # the HOG detector returns slightly larger rectangles than the real objects. 231 | # so we slightly shrink the rectangles to get a nicer output. 232 | pad_w, pad_h = int(0.15*w), int(0.05*h) 233 | cv2.rectangle(img, (x+pad_w, y+pad_h), (x+w-pad_w, y+h-pad_h), (0, 255, 0), thickness) 234 | 235 | def load_cascade_file(self, cascade_file='cascade.xml'): 236 | cascade_path = self.cascade_xml_dir + cascade_file 237 | self.logger.info("loading cascade file: " + cascade_path) 238 | self.cascade = cv2.CascadeClassifier(cascade_path) 239 | if self.cascade.empty(): 240 | Exception("cannot load cascade file: " + cascade_path) 241 | 242 | def detect(self, image_path): 243 | 244 | # read image 245 | self.logger.info("loading image file: " + image_path) 246 | img = cv2.imread(image_path) 247 | if img is None: 248 | Exception("cannot load image file: " + image_path) 249 | 250 | # detection 251 | self.logger.info("detecting") 252 | found = self.cascade.detectMultiScale(img, 1.1, 3) 253 | print found 254 | 255 | # join detection areas 256 | found_filtered = [] 257 | for ri, r in enumerate(found): 258 | for qi, q in enumerate(found): 259 | if ri != qi and self.inside(r, q): 260 | break 261 | else: 262 | found_filtered.append(r) 263 | 264 | # draw detection areas 265 | self.draw_detections(img, found) 266 | self.draw_detections(img, found_filtered, 3) 267 | print '%d (%d) found' % (len(found_filtered), len(found)) 268 | 269 | # show result 270 | # cv2.imshow('img', img) 271 | 272 | # write file 273 | head, tail = os.path.split(image_path) 274 | new_img_path = self.out_dir + tail + '.png' 275 | cv2.imwrite(new_img_path, img) 276 | 277 | # wait key 278 | ch = 0xFF & cv2.waitKey() 279 | if ch == 27: 280 | return 281 | cv2.destroyAllWindows() 282 | 283 | def detect_all(self): 284 | for file in self.test_img_files: 285 | self.detect(self.test_img_dir + file) 286 | 287 | def count_all_bboxes(self): 288 | sum = 0 289 | for annotation_file in self.my_annotation_files: 290 | annotation_path = self.my_annotation_dir + annotation_file 291 | # print annotation_path 292 | bboxes = my_util.my_unpickle(annotation_path) 293 | sum += len(bboxes) 294 | return sum 295 | 296 | def get_annotation_existence_list(self): 297 | return [file + '.pkl' in self.my_annotation_files for file in self.pos_img_files] 298 | 299 | def get_annotated_image_files(self): 300 | return [os.path.splitext(annotation_file)[0] for annotation_file in self.my_annotation_files] 301 | 302 | def get_annotation_path(self, img_file): 303 | annotation_path = self.my_annotation_dir + img_file + '.pkl' 304 | return annotation_path 305 | 306 | def get_img_file_by_annotation_file(self, annotation_file): 307 | img_file = os.path.splitext(annotation_file)[0] 308 | return img_file 309 | 310 | 311 | def load_annotation(self, img_file): 312 | bboxes = [] 313 | if img_file in self.get_annotated_image_files(): 314 | 315 | # annotation path 316 | annotation_path = self.get_annotation_path(img_file) 317 | 318 | self.logger.info('loading annotation file: %s', annotation_path) 319 | 320 | # load pickle 321 | bboxes = my_util.my_unpickle(annotation_path) 322 | return bboxes 323 | 324 | def save_annotation(self, img_file, bboxes): 325 | annotation_path = self.get_annotation_path(img_file) 326 | self.logger.info('saving annotation data: %s', annotation_path) 327 | my_util.my_pickle(bboxes, annotation_path) 328 | 329 | def draw_areas(self, im_orig, start_pt, end_pt, bboxes): 330 | im_copy = im_orig.copy() 331 | # draw rectangles 332 | if start_pt is not None and end_pt is not None: 333 | cv2.rectangle(im_copy, start_pt, end_pt, (0, 0, 255), 1) 334 | for box in bboxes: 335 | cv2.rectangle(im_copy, box[0], box[1], (0, 255, 0), 1) 336 | return im_copy 337 | 338 | def read_img(self, img_file): 339 | self.logger.info('loading image file: %s', img_file) 340 | img_path = self.pos_img_dir + img_file 341 | 342 | # read image 343 | cv_img = cv2.imread(img_path) 344 | return cv_img 345 | 346 | def draw_bounding_boxes_for_all(self): 347 | self.logger.info("begin drawing bounding boxes") 348 | for file_name in self.my_annotation_files: 349 | 350 | img_file = self.get_img_file_by_annotation_file(file_name) 351 | cv_img = self.read_img(img_file) 352 | bboxes = self.load_annotation(img_file) 353 | cv_bbox_img = self.draw_areas(cv_img, None, None, bboxes) 354 | # draw bounding box 355 | self.write_bbox_img(img_file, cv_bbox_img) 356 | 357 | 358 | def write_bbox_img(self, img_file, cv_bbox_img): 359 | bbox_path = self.my_annotation_img_dir + 'bbox_' + img_file 360 | self.logger.info('saving bounding box data: %s', bbox_path) 361 | cv2.imwrite(bbox_path, cv_bbox_img) 362 | 363 | def reload_my_annotation_files(self): 364 | self.my_annotation_files = [file_name for file_name in os.listdir(self.my_annotation_dir) if not file_name.startswith('.')] 365 | self.my_annotation_files.sort() 366 | 367 | if __name__ == '__main__': 368 | 369 | logging.root.setLevel(level=logging.INFO) 370 | 371 | dataset = ImageDataSet() 372 | 373 | dataset.create_crop_with_my_annotation() 374 | # dataset.create_samples(True, 24, 24) 375 | # dataset.train_cascade('HOG', 0.4, 0.995 24, 24) 376 | 377 | # dataset.load_cascade_file() 378 | # dataset.detect_all() 379 | # dataset.detect('./INRIAPerson/Train/pos/crop001509.png') 380 | # -------------------------------------------------------------------------------- /image_util.py: -------------------------------------------------------------------------------- 1 | #coding: utf-8 2 | 3 | import cv2 4 | import pickle 5 | 6 | def show_cv_image(window_name, img): 7 | window_name_encoded = unicode(window_name, 'utf-8').encode('cp932') 8 | cv2.namedWindow(window_name_encoded) 9 | cv2.imshow(window_name, img) 10 | cv2.waitKey(0) 11 | cv2.destroyAllWindows() 12 | 13 | def image_crop(img,left_top, right_bottom): 14 | return img[left_top[1]:right_bottom[1], left_top[0]:right_bottom[0]] 15 | -------------------------------------------------------------------------------- /inria_person_dataset.py: -------------------------------------------------------------------------------- 1 | #coding: utf-8 2 | 3 | import os 4 | import re 5 | import cv2 6 | import image_util as iu 7 | import logging 8 | from image_dataset import ImageDataSet 9 | 10 | class InriaPersonDataSet(ImageDataSet): 11 | 12 | BBOX_DIR = 'bounding_box' 13 | ANNOTATION_DIR = 'annoation' 14 | CONFIG_YAML = 'inria_config.yml' 15 | 16 | def __init__(self): 17 | 18 | ImageDataSet.__init__(self) 19 | 20 | # set dataset path 21 | self.annotation_dir = self.config['dataset']['annotation_dir'] 22 | 23 | # set output path 24 | self.bounding_box_out_dir = self.config['output']['bounding_box_out_dir'] 25 | 26 | # create output paths 27 | if not os.path.isdir(self.bounding_box_out_dir): 28 | os.makedirs(self.bounding_box_out_dir) 29 | 30 | 31 | def parse_annotation_file(self, img_file_name): 32 | 33 | # image annotation path 34 | annotation_path = self.annotation_dir + os.path.splitext(img_file_name)[0] + '.txt' 35 | 36 | # open annotation file 37 | f = open(annotation_path) 38 | lines = f.readlines() 39 | f.close() 40 | 41 | # parse annotation file 42 | object_list = [] 43 | object_info = {} 44 | ground_truth = None 45 | img_size = None 46 | 47 | for line in lines: 48 | # print line, 49 | 50 | # get image size 51 | m = re.match(r'Image size \(X x Y x C\) : (\d+) x (\d+) x 3', line) 52 | if m: 53 | img_size = (int(m.group(1)), int(m.group(2))) 54 | # print img_size 55 | 56 | # get ground truth 57 | m = re.match(r'Objects with ground truth : (\d+)', line) 58 | if m: 59 | ground_truth = int(m.group(1)) 60 | # print ground_truth 61 | 62 | if line.find('# Details for object') != -1: 63 | object_info = {} 64 | # print '# Details for object' 65 | 66 | # get center 67 | m = re.match(r'Center point on object (\d)+ "PASperson" \(X, Y\) : \((\d+), (\d+)\)', line) 68 | if m: 69 | center = (int(m.group(2)), int(m.group(3))) 70 | # print center 71 | object_info['center'] = center 72 | 73 | # get bounding box 74 | m = re.match(r'Bounding box for object (\d+) "PASperson" \(Xmin, Ymin\) - \(Xmax, Ymax\) : \((\d+), (\d+)\) - \((\d+), (\d+)\)', line) 75 | if m: 76 | bounding_box = [(int(m.group(2)), int(m.group(3))), (int(m.group(4)), int(m.group(5)))] 77 | # print bounding_box 78 | object_info['bounding_box'] = bounding_box 79 | object_list.append(object_info) 80 | 81 | # check number of objects 82 | if len(object_list) != ground_truth: 83 | Exception("parsing error: ground truth does not match with object number.") 84 | return None 85 | 86 | # create annotation info 87 | annotation_info = { 88 | 'img_size': img_size, 89 | 'ground_truth': ground_truth, 90 | 'object_list': object_list 91 | } 92 | 93 | return annotation_info 94 | 95 | def draw_bounding_boxes_for_all(self): 96 | self.logger.info("begin drawing bounding boxes") 97 | for file_name in self.pos_img_files: 98 | 99 | # draw bounding box 100 | self.draw_bounding_boxes_and_write_file(file_name) 101 | 102 | def draw_bounding_boxes_and_write_file(self, file_name): 103 | 104 | file_path = self.pos_img_dir + file_name 105 | self.logger.info("drawing bounding box: " + file_path) 106 | 107 | # read image 108 | img = cv2.imread(file_path, cv2.IMREAD_COLOR) 109 | 110 | # read annotation file to get annotation info 111 | annotation_info = self.parse_annotation_file(file_name) 112 | 113 | # iterate object list and draw bounding boxes 114 | for object_info in annotation_info['object_list']: 115 | bounding_box = object_info['bounding_box'] 116 | cv2.rectangle(img, bounding_box[0], bounding_box[1], (0, 0, 255), 5) 117 | 118 | # output file 119 | cv2.imwrite(self.bounding_box_out_dir + 'b_' + file_name, img) 120 | 121 | def create_crop_for_all(self): 122 | self.logger.info("begin creating crop image") 123 | for file_name in self.pos_img_files: 124 | 125 | # create crop 126 | self.create_crop_write_file(file_name) 127 | 128 | def create_crop_write_file(self, file_name): 129 | 130 | file_path = self.pos_img_dir + file_name 131 | self.logger.info("creating crop image: " + file_path) 132 | 133 | # read image 134 | img = cv2.imread(file_path, cv2.IMREAD_COLOR) 135 | 136 | # read annotation file to get annotation info 137 | annotation_info = self.parse_annotation_file(file_name) 138 | 139 | # iterate object list and create crop images 140 | for i, object_info in enumerate(annotation_info['object_list']): 141 | crop_box = object_info['bounding_box'] 142 | cropped_img = iu.image_crop(img, crop_box[0], crop_box[1]) 143 | out_file_name = 'c_' + os.path.splitext(file_name)[0] + '_' + str(i) + '.' + os.path.splitext(file_name)[1] 144 | cv2.imwrite(self.cropped_dir + out_file_name, cropped_img) 145 | 146 | def create_positive_dat_by_image_size(self): 147 | output_text = "" 148 | self.logger.info("begin creating positive.dat") 149 | for file_name in self.pos_img_files: 150 | 151 | file_path = self.pos_img_dir + file_name 152 | annotation_info = self.parse_annotation_file(file_name) 153 | output_text += "%s %d " % (file_path, annotation_info['ground_truth']) 154 | for object_info in annotation_info['object_list']: 155 | x_min, y_min = object_info['bounding_box'][0] 156 | x_max, y_max = object_info['bounding_box'][1] 157 | w = x_max - x_min 158 | h = y_max - y_min 159 | output_text += "%d %d %d %d " % (x_min, y_min, w, h) 160 | output_text += "\n" 161 | # print output_text 162 | self.logger.info("writing data to positive.dat") 163 | f = open('positive.dat', 'w') 164 | f.write(output_text) 165 | f.close() 166 | self.logger.info("completed writing data to positive.dat") 167 | 168 | if __name__ == '__main__': 169 | 170 | logging.root.setLevel(level=logging.INFO) 171 | 172 | inria = InriaPersonDataSet() 173 | 174 | # inria.draw_bounding_boxes_for_all() 175 | 176 | # inria.create_crop_for_all() 177 | inria.create_samples(40, 40) 178 | inria.train_cascade('HOG', 0.4, 0.995, 40, 40) 179 | 180 | inria.load_cascade_file() 181 | inria.detect_all() 182 | # inria.detect('./INRIAPerson/Train/pos/crop001509.png') 183 | # -------------------------------------------------------------------------------- /my_util.py: -------------------------------------------------------------------------------- 1 | try: 2 | import cPickle as pickle 3 | except: 4 | import pickle 5 | 6 | 7 | def my_pickle(data, path): 8 | f = open(path, 'wb') 9 | pickle.dump(data, f) 10 | f.close() 11 | 12 | def my_unpickle(path): 13 | f = open(path, 'rb') 14 | data = pickle.load(f) 15 | f.close() 16 | return data 17 | 18 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | from image_dataset import ImageDataSet 5 | import logging 6 | 7 | if __name__ == '__main__': 8 | 9 | logging.root.setLevel(level=logging.INFO) 10 | dataset = ImageDataSet() 11 | dataset.train_cascade(feature_type='HOG', max_false_alarm_rate=0.4, min_hit_rate=0.995, width=24, height=24, pos_rate=0.5) 12 | 13 | --------------------------------------------------------------------------------