Inspiration
As a final project for my neural networks class, my group partnered with a company which specializes in protecting peoples' online reputations. The company has a dataset of images containing objects you wouldn't want tagged on your Facebook profile (i.e. drugs, weapons, etc.). They currently have a system in place that is supposed to detect these objects in your personal images. We were looking to see if we could improve it.
What it does
My team and I customized and trained a neural network to detect these objects in images, flagging them for review. We achieved high model performance while reducing model training time from 3 weeks to 2 days and reducing inference time and memory footprint.
How I built it
The dataset had a heavy class imbalance (many more images with beer bottles than weapons, for instance). We wanted to find a network capable of handling this imbalance. We researched and settled on using a model built with a RetinaNet backbone. RetinaNet uses a focal loss function, which helps to alleviate this class imbalance problem.
It turned out that Facebook AI Research labs had a very new (and still rough around the edges) framework called Detectron, build specifically for detecting objects in images. We decided to implement our model within the Detectron framework, and leverage transfer learning techniques to aid in training our model to the specific task at hand.
Challenges I ran into
Setting up Detectron was no small feat. Additionally, the dataset we used was not in the proper format, which required some manipulation.
Accomplishments that I'm proud of
Successfully working with one of the latest and greatest object detection models to achieve a specific custom goal.
Building a testing harness compatible with any model that produces results in MSCOCO format to view and compare both quantitative (i.e. statistical) and qualitative (visualized detections) results across multiple models.
Building the foundations for a continuous data integration pipeline to translate the company's SQL data into the standard MSCOCO JSON format for sustained use with future model architectures.
What I learned
Anything is possible.
What's next for this project
This was a class project, which we presented to the partner company. They liked it so much they are looking to contract us for hire to continue developing the project.
Built With
- amazon-web-services
- caffe2
- detectron
- mscoco
- python
Log in or sign up for Devpost to join the conversation.