How to Train YOLOv5-Classification on a Custom Dataset

Key Takeaways

YOLO you only look once arguably the most influential and successful image classification model ever wait what I mean yeah according to the latest release notes they are doing classification now so in our last video we are talking about yellow V5 being capable of commit segmentation and today we'll be talking about image classification by the way link to that video will be in the description below and in the card displayed right now but before we start we actually need to explain why it's such a big deal so YOLO has algorithm and YOLO V5 as a framework are very good at training object detection models yeah object detection not image classification those are definitely two different things object detection is where you say what object do you see and where on the image you see that so for example over here we say that we see raccoon and we precisely where the raccoon is located image classification is simpler task we just say I see raccoon on the image why would I train image classification model if I can train object detection there are actually multiple reasons where you would do that for starters image classification models are lighter they simply consume less space in the memory second of all you can train a model with significantly more classes than in case of object detection typically hundreds or even thousands not to mention that the actual inference takes less time which is very important in many scenarios every architecture has pros and cons but you can definitely come up with scenario where the image classification model is the desired one roboflow has actually a very good article about different computer vision tasks once again the link is in the description below it's time to get our hands dirty and train our first YOLO V5 classification model so as previously with image segmentation we have a dedicated blog post that describe the whole process but we will go straight into the Google collab that is linked over here Google Cloud app opens a new tab now we can just save copy on our hard drive the next steps are similar like in the case of object detection and image segmentation so we just connect our jupyter notebook to the python environment with gpus and after that is done we run the first cell the first cell is mostly responsible for colliding the repository and setting up the python project installing the dependencies mainly pytorch now we can make sure that everything installed properly so we will download the weights of every model that is accessible within a yellow V5 repository and use one of those models to infer on example image so the first cell is actually downloading an example image the second one loads model into the memory and produce classification the name of the image slightly ruins the surprise bananas JPEG but yeah our model is 96 sure that there are bananas on the image so now the time has finally come and we'll actually train the model ourselves on our custom data set before we start however I guess it's worth to mention that the format of data for image classification is different than usual yellow format for object detection typically when you train your model you work with something called yellow darknet txt format where every Row in the txt file is separate detection and you have the ID of the class as well as the coordinates of the bounding box in it for image classification we don't have any additional annotation file so how do we know what is the class of the object visible on the image we know that by the name of the directory that the image is located so for example over here inside the train data set we have class 1 directory and we know that all images in that directory are of class 1 and similarly over here we know that all images in that directory are of class number two now the moment that we've all been waiting for let's train our own classification model so we will use tomato classification data set that is the one that we used in our tutorial so let's follow the link in the tutorial we click download select the format so in our case like I said it will be file structure click continue and copy the snippet straight into our jupyter notebook over here in our Jupiter notebook we have few cells one of them will create the data sets directory the other one will download the data set from roboflow website protect your API key at all costs don't share it with anybody our data set has downloaded correctly and we can start the training there are just few more cells that we need to run the first one is setting the data set name as an environment variable and then we are ready to run the actual training script we use 100 epochs over here so that may take a little bit of time to train although it's nothing compared to image segmentation or object detection actually it wasn't so bad took less than 15 minutes finally we are ready to validate our model so let's use script classified slash valve slash py to run our validation on pretty small but still validation set and we are 100 correct and that's it leave a comment under the video if you like it we have a lot of similar content in the pipeline mainly about yellow V5 yellow V6 and yellow V7 leave a sap if you find this video useful and my name was Peter and I see you next time bye