Object detection on colors in a Rubik's Cube
git clone [email protected]:CubeCV/color-detector.git
cd color-detector
conda env create -f env.yml
conda activate color-detector-envCUDA_VISIBLE_DEVICES=<GPU ID> python main.pyCUDA_VISIBLE_DEVICES=<GPU ID> python validate.py --model=path/to/best.ptThe path to best.pt is currently models/best.pt.
CUDA_VISIBLE_DEVICES=<GPU ID> python predict.py --model=path/to/best.pt --image=path/to/imageThe path to best.pt is currently models/best.pt.
The data for this project is taken from https://universe.roboflow.com/psst/mofang-i7ha2.
We have split it into train (1198 images + labels), val (100 images + labels) , and test (100 images + labels) in the datasets directory. There is also a datasets/custom directory which contains images not from the aforementioned dataset.
We used the pretrained yolov8n.pt model to build our own color detector model. Our custom model is models/best.pt.
Our model performs well when the cube is isolated (no background distractions) as can be seen in runs/detect/predict/. Hence, we would need to segment the cube prior to running the color detection model, which is phase 1 of the CubeCV project.