af-a / dataset_interface

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

dataset_interface

This repository includes:

  • A common interface for interacting with different datasets, in context of the b-it-bots RoboCup teams.
  • Tools/pipelines for automatically segmenting object masks using the green box (picture below), using these masks to synthesize data for training object detection models, and training on the synthesized data

Note: tested only with Python 3.

Installation

Since the data API from pycocotools requires building Cython modules, pip install -e . and python setup.py develop does not seem to work. Arch users may have to install the tk windowing toolkit manually as a system dependency.

python setup.py install --user

Common interface for datasets

A more detailed description of how this interface works can be found in the dataset interface documentation

A sample config file for the COCO dataset: sample_coco_configs.yml

from dataset_interface.coco import COCODataAPI

data_dir = 'data/dir'
config_file = 'config/file.yml'
coco_api = COCODataAPI(data_dir, config_file)

# return a dictionary mapping image ID to ImageInfo objects
images = coco_api.get_images_in_category('indoor')
first_image = next(iter(images.values()))
print(first_image.image_path)   # image file location
print(first_image.url)          # image URL for download
print(first_image.id)           # image unique ID in the dataset

# return a dictionary containing all categories under the 'indoor' super category,
# mapping category ID's to Category objects
indoor_cats = coco_api.get_sub_categories('indoor')
# mapping category ID's to category names
indoor_cat_names = coco_api.get_sub_category_names('indoor')

Object Detection and Classification Model Training

We primarily train detection models using the tools and model definitions from torchvision as described in the Torchvision Object Detection Finetuning Tutorial. Guidelines for training and using models can be found in docs/object_detection.md

Image Comparison Model Training

We use Siamese networks for object/person comparison and subsequent recognition. Instructions for training and using a comparison model are given in docs/image_comparison.md.

Image augmentation

We aim to ease the process of generating data for object detection. Using the green box in the picture below, it's possible to automatically segment objects and transform them onto new backgrounds to create new training data for an object detection model. A more detailed documentation of how we solve this problem can be found in docs/image_augmentation.md.

Green Box

About

License:GNU General Public License v3.0


Languages

Language:Jupyter Notebook 98.5%Language:Python 1.3%Language:C++ 0.1%Language:C 0.1%Language:Makefile 0.0%