This is a simple Python script using YOLOv5 to perform object detection on custom classes. The script is designed to detect the following object classes in an input image:
- Red spot
- Black spot
- Nipple
- Extraction point
- Point flame
The script uses PyTorch, OpenCV, and NumPy libraries for processing and visualisation.
To run the script, you need to have the following libraries installed:
- PyTorch
- OpenCV
- NumPy
You can install them using pip: pip install torch torchvision opencv-python numpy
- Download the pre-trained YOLOv5 model weights file
weight_all.pt
and place it in the same directory as the script. - Place an input image (e.g.,
IMG_3245.jpg
) in theResources\Full_image_set_resized
directory. - Run the script:
python your_script_name.py
The script will process the input image, detect the objects, and display the results with bounding boxes and class labels.
You can modify the following parameters in the script to suit your needs:
weights
: Path to the pre-trained YOLOv5 model weights file.device
: Set to'cpu'
to use the CPU, or'cuda:0'
to use the GPU.classes
: List of class names to detect.conf_threshold
: Confidence threshold for object detection (default is 0.4).nms_threshold
: Non-maximum suppression threshold for merging overlapping bounding boxes (default is 0.5).
If you encounter any issues while running the script, please ensure that:
- All required libraries are installed correctly.
- The pre-trained YOLOv5 model weights file (
weight_all.pt
) is in the same directory as the script. - The input image is placed in the
Resources\Full_image_set_resized
directory. - The input image file path is correct and the image file is not corrupted.
If you still have issues, consider checking the following resources for further assistance:
- YOLOv5 GitHub repository: The official repository for YOLOv5, which contains the latest updates, documentation, and examples.
- PyTorch forum: A community-driven forum for discussing and resolving issues related to PyTorch.
- OpenCV forum: A community-driven forum for discussing and resolving issues related to OpenCV.
- NumPy mailing list: A mailing list for NumPy users