It's a project to detect traffic cones and recognize the colors as well. I used yolov5 to train and detect cones. Furthermore, I used k-means to determine the dominant color to classify cone color. Currently, the supported colors are red, yellow, green, and blue. Other colors are classified as unknown.
I used a self-collected cone dataset with 303 cone images. It's not a perfect practice because it's a small dataset. I also need to annotate the images myself. Here, I utilized an online annotation website Roboflow, it provides services such as annotation, pre-processig, and augmentation. However, it has limitation of 1,000 source images and 5,000 generated images for free users.
Model
├── cone detection: yolov5s
└── color recognition: dominant color (k-means)You can try the codes in colab if you are interested in.
If you have an annotated dataset, you can train directly use train.ipynb
If you want to detect cones directly, use predict.ipynb
You should use the weights I trained in model. Besides, I customized some files of yolov5, which are located in utils folder, you need to use them as well.
I clipped a video from one research project of ETH Zurich to test the peroformance.
