Last Updated on 22/04/2026 by Eran Feit
Detectron2 custom dataset training means taking your own images (not COCO), labeling them with polygon masks, registering them in Detectron2, and fine-tuning Mask R-CNN so it can detect and segment your specific objects.
In this tutorial, we’ll walk through a complete pipeline using a fruit dataset: apples, bananas, strawberries, grapes, oranges, and lemons.register_coco_instances, and confirm the masks render correctly.
Finally, we’ll load the trained weights, run inference on fresh test images, and visualize the predictions with colored instance masks.
Along the way, we’re going to reuse high-value concepts like DatasetCatalog, MetadataCatalog, COCO annotations, and the Detectron2 model zoo.
Detectron2 Architecure Introduction: training a custom fruit detector with Detectron2 the easy way This guide shows how to build a complete Detectron2 instance segmentation pipeline for fruit images, using Mask R-CNN to detect apples, bananas, grapes, strawberries, oranges, and lemons — and draw masks around them, not just boxes.
See how I run real-time object detection on live camera input, check out my Jetson Nano walkthrough: https://eranfeit.net/how-to-classify-objects-in-live-camera-using-jetson-nano/
Link for the video tutorial : https://youtu.be/JbEy4Eefy0Y
You can download the code here : https://eranfeit.lemonsqueezy.com/buy/4e044baf-d59e-46f4-bf77-c9021c6d07d6
or here : https://ko-fi.com/s/2119cfd494
Here is a link to the post for Medium users .
Master Computer Vision
Follow my latest tutorials and AI insights on my
Personal Blog .
Beginner Complete CV Bootcamp
Foundation using PyTorch & TensorFlow.
Get Started → Interactive Deep Learning with PyTorch
Hands-on practice in an interactive environment.
Start Learning → Advanced Modern CV: GPT & OpenCV4
Vision GPT and production-ready models.
Go Advanced → Getting our fruit dataset labeled in COCO format A Detectron2 instance segmentation project always begins with data.
Before we can train Detectron2 instance segmentation, we need high-quality labels.
We export the labels to a COCO-style JSON file.
he important part for “train detectron2 custom dataset” is consistency:
We keep a third split, Test, with no labels.
Below is the prep workflow in a code-style block so you can copy it into your project notes or README.
Email me for the sample dataset .
### Prepare COCO-style annotations using an online annotation tool. # Coco format annotation using an online annotation tool ### The dataset images are fruit photos (apples, bananas, etc.) that we collected. # dataset from here : images from Google : # Folder under Detectron2 :-> Train-custom-Object-Detection-model\Fruits_for_detectron2 ### These are the classes we want Detectron2 to learn. Labels : Apple , Strawberry, Orange, Grapes, Banana , Lemon ### Open the annotation tool in your browser. # open the annotation tool website ### Upload all training images so we can draw polygons around each fruit. # upload the train images ### Choose "Object detection" / "Instance segmentation" mode. # choose Object detection ### For each object, draw a polygon around the fruit to create a precise mask. # select polygon on the right side bottom for each item in each image ### When you finish labeling, export the annotations in COCO format. # choose "actions" then "export annotions" and choose Coco format ### Copy the exported JSON file into the Train folder next to the images. # copy the Json file to the train images folder ### Repeat labeling for the validation split so we can evaluate during training. # Repeat : # Do the same with the validate images !!! ### We do not need labels for the test split because we only want predictions later. # The test images folder can stay without annotations Summary of this part:
We defined our fruit categories for Detectron2 instance segmentation. We labeled Train and Validate with polygons and exported COCO JSON per split. We kept an unlabeled Test folder for inference. Clean, consistent labels at this step are the foundation for every good Mask R-CNN model. If you’re interested in classic image classification (not masks), I also compare deep models in this guide: https://eranfeit.net/tensorflow-image-classification-tutorial-resnet50-vs-mobilenet /
Loading the fruit dataset into Detectron2 and previewing it Now that we have COCO-style JSON files, we register them with Detectron2.
Detectron2 exposes two helpers that make this easy: register_coco_instances and MetadataCatalog.register_coco_instances connects a dataset name to an image directory and its COCO annotations.MetadataCatalog then stores extra info like class names so the visualizer can draw correct labels.
We load my_dataset_train and my_dataset_val and immediately query them using DatasetCatalog.get().
We use OpenCV (cv2.imread) to read the first training image.Visualizer to overlay the masks and boxes from the annotations onto the image.
This step is part of “custom object detection with detectron2,” but we are actually doing instance segmentation with Mask R-CNN.
Here’s the full data loading and visualization block.### explanation line immediately above it.
### Import Detectron2 so we can work with instance segmentation and training utilities. import detectron2 ### Import OpenCV for image reading and window display. import cv2 ### Import the Visualizer to draw annotations (masks, boxes, class labels). from detectron2 . utils . visualizer import Visualizer ### Import catalogs that store dataset metadata and items. from detectron2 . data import MetadataCatalog , DatasetCatalog ### Import the helper that registers COCO-style datasets. from detectron2 . data . datasets import register_coco_instances ### Register our datasets so Detectron2 knows where images and COCO JSON live. # register datasets ### Register the training dataset with a name, empty metadata dict, the JSON file, and the images folder. register_coco_instances ( " my_dataset_train " , {}, " Train-custom-Object-Detection-model/Fruits_for_detectron2/Train/labels_my-project-name_2023-12-04-07-26-09.json " , " Train-custom-Object-Detection-model/Fruits_for_detectron2/Train " ) ### Register the validation dataset the same way. register_coco_instances ( " my_dataset_val " , {}, " Train-custom-Object-Detection-model/Fruits_for_detectron2/Validate/labels_my-project-name_2023-12-04-07-39-25.json " , " Train-custom-Object-Detection-model/Fruits_for_detectron2/Validate " ) ### Grab metadata for the train split (class names, colors, etc.). train_metedata = MetadataCatalog . get ( " my_dataset_train " ) ### Grab the list of training image dictionaries. train_datasets_dicts = DatasetCatalog . get ( " my_dataset_train " ) ### Grab metadata for the validation split. val_metedata = MetadataCatalog . get ( " my_dataset_val " ) ### Grab the list of validation image dictionaries. val_datasets_dicts = DatasetCatalog . get ( " my_dataset_val " ) ### Look at the first training sample to understand its structure. first_dict = train_datasets_dicts [ 0 ] ### Print that first sample info (file path, annotations, size). print ( first_dict ) ### Extract useful fields from the first training record. file_name = first_dict [ ' file_name ' ] height = first_dict [ ' height ' ] width = first_dict [ ' width ' ] image_id = first_dict [ ' image_id ' ] annotations = first_dict [ ' annotations ' ] ### Read the actual image using OpenCV. img = cv2 . imread ( file_name ) ### Create a Visualizer object, telling it to use our training metadata. visual = Visualizer ( img [:, :, :: - 1 ] , metadata = train_metedata , scale = 0.5 ) ### Ask the Visualizer to draw the ground-truth annotations on top of the image. vis = visual . draw_dataset_dict ( first_dict ) ### Convert the visualized output to an image array we can view. img2 = vis . get_image () ### Convert BGR to RGB for correct color when displaying in some viewers. img_rgb = cv2 . cvtColor ( img2 , cv2 . COLOR_BGR2RGB ) ### Show the preview window so we can confirm masks and boxes look correct. cv2 . imshow ( " img_rgb " , img_rgb ) ### Wait for a key press before closing the preview window. cv2 . waitKey ( 0 ) Summary of this part:
We used register_coco_instances to make Detectron2 aware of our dataset. We previewed one training image with true masks and boxes using Visualizer. This confirms our COCO JSON is valid and our fruit classes map correctly. Visual sanity checks now save hours later when training. Training the model on Windows (CPU-friendly version) Now we actually train Detectron2 on our fruit dataset.
We define paths, dataset names, and an output directory (My-Train-Detectron2).num_classes = 6 to match Apple, Strawberry, Orange, Grapes, Banana, Lemon.
We build a config object using get_cfg() and model_zoo.get_config_file(...).
Which dataset to train on Which dataset to test/evaluate on Batch size (IMS_PER_BATCH) Learning rate (BASE_LR) Training length (MAX_ITER) Number of workers for the data loader We set cfg.MODEL.DEVICE = "cpu" so this script can run on a normal Windows machine with no GPU.
We save the config to a pickle file so we can reload it later for inference.DefaultTrainer, which handles the training loop for us.
Here’s the full Windows training script, with explanations (###) above every instruction:
### Import dataset registration helper. from detectron2 . data . datasets import register_coco_instances ### Import DefaultPredictor (for inference) and DefaultTrainer (for training loops). from detectron2 . engine import DefaultPredictor , DefaultTrainer ### Import OS tools for paths and pickle for saving config. import os import pickle ### We'll use a model definition from the Detectron2 model zoo. # go to the model zoo to pick a config name that matches instance segmentation ### Set which config file we'll use (Mask R-CNN with ResNet-50 and FPN). config_file_path = " COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml " ### Name of the model we want weights from. model_name = " COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml " ### Where to save training outputs (weights, logs, config). output_dir = " My-Train-Detectron2 " ### We have 6 fruit classes (Apple, Strawberry, Orange, Grapes, Banana, Lemon). num_classes = 6 ### On Windows we train on CPU for simplicity. device = " cpu " ### Give friendly names to our datasets for training and validation. # Train train_dataset_name = " LP_train " ### Path to training images. train_images_path = r " Train-custom-Object-Detection-model/Fruits_for_detectron2/Train " ### Path to COCO annotations for training. train_json_annot_path = r " Train-custom-Object-Detection-model/Fruits_for_detectron2/Train/labels_my-project-name_2023-12-04-07-26-09.json " ### Validation dataset details. # Validate val_dataset_name = " LP_Test " ### Path to validation images. val_images_path = r " Train-custom-Object-Detection-model/Fruits_for_detectron2/Validate " ### Path to COCO annotations for validation. val_json_annot_path = r " Train-custom-Object-Detection-model/Fruits_for_detectron2/Validate/labels_my-project-name_2023-12-04-07-39-25.json " ### Register the training dataset so Detectron2 can load images and annotations. # Register the dataset #Register the train: register_coco_instances ( name = train_dataset_name , metadata = {}, json_file = train_json_annot_path , image_root = train_images_path ) ### Register the validation dataset. #Register the Validation: register_coco_instances ( name = val_dataset_name , metadata = {}, json_file = val_json_annot_path , image_root = val_images_path ) ### Import configuration utilities from Detectron2. from detectron2 . config import get_cfg ### Import model_zoo to load base configs and pretrained weights. from detectron2 import model_zoo ### Define a helper that builds and customizes the Detectron2 config. def get_train_cfg ( a_config_file_path , a_model_name , a_train_dataset_name , a_test_dataset_name , a_num_classes , device , output_dir ): ### Start from a default config object. cfg = get_cfg () ### Load a baseline Mask R-CNN config from the model zoo. cfg . merge_from_file ( model_zoo . get_config_file ( a_config_file_path )) ### Use pretrained weights for faster convergence. cfg . MODEL . WEIGHTS = model_zoo . get_checkpoint_url ( a_model_name ) ### Tell Detectron2 which dataset names correspond to training and testing. cfg . DATASETS . TRAIN = ( a_train_dataset_name , ) cfg . DATASETS . TEST = ( a_test_dataset_name , ) ### Number of data loader workers (higher = faster if you have CPU cores). cfg . DATALOADER . NUM_WORKERS = 2 ### How many images per iteration (batch size at the solver level). cfg . SOLVER . IMS_PER_BATCH = 2 # how many images per batch ### Base learning rate for the optimizer. cfg . SOLVER . BASE_LR = 0.00025 # learning rate ### Training length in iterations (similar idea to epochs). cfg . SOLVER . MAX_ITER = 1000 # how many iterations (like epochs) ### When to drop the learning rate. Empty list = don't schedule drops. cfg . SOLVER . STEPS = [] # this will reduce the learning rate during the training ### Number of target classes (your custom fruit classes). cfg . MODEL . ROI_HEADS . NUM_CLASSES = a_num_classes ### Run on CPU or GPU. cfg . MODEL . DEVICE = device ### Where to save checkpoints and logs. cfg . OUTPUT_DIR = output_dir ### Return the configured training object. return cfg ### main entry point for training on Windows. def main () : ### Make sure the output directory exists. os . makedirs ( output_dir , exist_ok =True ) ### Build the training config with our dataset names and hyperparameters. cfg = get_train_cfg ( config_file_path , model_name , train_dataset_name , val_dataset_name , num_classes , device , output_dir ) ### Path where we will save the config object for later inference. cfg_save_path = " My-Train-Detectron2/IS_cfg.pickle " # IS -> Instance segmentation ### Save the config using pickle so we can reload it later. with open ( cfg_save_path , ' wb ' ) as f : pickle . dump ( cfg , f , protocol = pickle . HIGHEST_PROTOCOL ) # save the cfg ### Create a DefaultTrainer with our config. trainer = DefaultTrainer ( cfg ) ### Start training from scratch (resume=False). trainer . resume_or_load ( resume =False ) ### Launch the training loop. trainer . train () ### Only run main() if this file is executed directly. if __name__ == ' __main__ ' : main () Summary of this part:
We built a Detectron2 config for Mask R-CNN instance segmentation. We trained on CPU for clarity, which is great for demos and “it actually runs on my laptop.” We saved the config to IS_cfg.pickle, which we’ll reload later for inference. This section directly supports the keyword “train detectron2 custom dataset.” For another segmentation approach (medical polyps instead of fruit), you can read my U-Net tutorial here: https://eranfeit.net/u-net-medical-segmentation-with-tensorflow-and-keras-polyp-segmentation/
Training on Ubuntu / WSL with GPU acceleration Training Detectron2 instance segmentation on CPU is fine for testing, but serious training is faster on GPU.device = "cuda" and a longer training schedule.
We first create and activate a Conda environment, install PyTorch with CUDA support, and build Detectron2 from source.
We again register the dataset splits (LP_train, LP_test) and point to the COCO JSON and images folders.Fruits_for_detectron2/Train./mnt/c/..., but here we assume the dataset is local in the Ubuntu environment.
We generate a config using get_train_cfg, similar to Windows, but with two important differences:
cfg.MODEL.DEVICE = "cuda" so Detectron2 uses the GPU.cfg.SOLVER.MAX_ITER = 3500, which means we train longer and (usually) get better accuracy. We save the config pickle (IS_cfg-Ubunto.pickle) and start training with DefaultTrainer.
Here’s the full Ubuntu / GPU training script in one block with explanations above each line:
### Notes for WSL / Ubuntu environment setup: #WSL - Lynux # create WSL Ubunto enviroment ### Activate Conda and create a clean env for Detectron2 + CUDA. # open WSL in c: #conda create -n detectorn99 python=3.9 #conda activate detectorn99 ### Install a CUDA-enabled PyTorch build. #install Pytorch #conda install pytorch==1.13.0 torchvision==0.14.0 torchaudio==0.13.0 pytorch-cuda=11.6 -c pytorch -c nvidia ### Clone Detectron2, enter its folder, and build it in develop mode. #cd detectron2 # inside the git clone folder #cd Train Folder #python setup.py build develop ### Run the actual training script. # run the python train code: #python Step4-Train-The-Model-Ubunto-WSL.py ### Import dataset registration and trainer utilities. # register the dataset from detectron2 . data . datasets import register_coco_instances from detectron2 . engine import DefaultTrainer ### Import OS and pickle to save config objects. import os import pickle ### We'll still use a Mask R-CNN config from the Detectron2 model zoo (instance segmentation capable). # you can see many models in the table : "Instance segmentation" ### Define which config and model weights to start from. # object detection config_file_path = " COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml " model_name = " COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml " ### Output directory for Ubuntu/GPU run. output_dir = r " My-Train-DetectRon2-Ubuntu " ### Our dataset still has 6 fruit classes. num_classes = 6 ### We want to train on GPU here. device = " cuda " ### Training dataset name and paths. #train train_dataset_name = " LP_train " ### Example paths for images and annotations when running inside Ubuntu/WSL. #train_images_path = r"/mnt/c/Python-Code/Best-Object-Detection-models/Detectron2/Fruits_for_detectron2/Train" #train_json_annot_path = r"/mnt/c/Python-Code/Best-Object-Detection-models/Detectron2/Fruits_for_detectron2/Train/labels_my-project-name_2023-12-04-07-26-09.json" train_images_path = r " Fruits_for_detectron2/Train " train_json_annot_path = r " Fruits_for_detectron2/Train/labels_my-project-name_2023-12-04-07-26-09.json " ### Validation dataset name and paths. #Validation val_dataset_name = " LP_test " val_images_path = r " Fruits_for_detectron2/Validate " val_json_annot_path = r " Fruits_for_detectron2/Validate/labels_my-project-name_2023-12-04-07-39-25.json " ### Actually register both datasets with Detectron2. # register the dataset: # ===================== ### Register the train split. # register the train register_coco_instances ( name = train_dataset_name , metadata = {}, json_file = train_json_annot_path , image_root = train_images_path ) ### Register the validation/test split. # register the test register_coco_instances ( name = val_dataset_name , metadata = {}, json_file = val_json_annot_path , image_root = val_images_path ) ### Import config helpers and model zoo from Detectron2. from detectron2 . config import get_cfg from detectron2 import model_zoo ### Build a training config specialized for GPU training and longer runs. def get_train_cfg ( a_config_file_path , a_model_name , a_train_dataset_name , a_test_dataset_name , a_num_classes , device , output_dir ) : ### Start with a base Detectron2 config. cfg = get_cfg () ### Merge in a known-good Mask R-CNN config from the model zoo. cfg . merge_from_file ( model_zoo . get_config_file ( a_config_file_path )) # get the "config_file_path" ### Use pretrained weights so training converges faster. cfg . MODEL . WEIGHTS = model_zoo . get_checkpoint_url ( a_model_name ) #model_zoo.get_checkpoint_url(a_model_name) ### Tell Detectron2 which datasets are train and test. cfg . DATASETS . TRAIN = ( a_train_dataset_name , ) cfg . DATASETS . TEST = ( a_test_dataset_name , ) ### Number of worker processes for data loading. # number of workers -> how to stress the CPU cfg . DATALOADER . NUM_WORKERS = 2 ### Images per batch (higher uses more VRAM). cfg . SOLVER . IMS_PER_BATCH = 2 # how many images per batch ### Learning rate for optimization. cfg . SOLVER . BASE_LR = 0.00025 # learning rate ### Total iterations (similar to epochs). Longer than Windows run. cfg . SOLVER . MAX_ITER = 3500 # how many iterations (like epohcs) !!!!!!!!!!!! ### Keep learning rate constant instead of stepping it down. cfg . SOLVER . STEPS = [] # this will not reduce the learning rate during training ### Tell the ROI heads how many classes we expect. cfg . MODEL . ROI_HEADS . NUM_CLASSES = a_num_classes ### Use CUDA/GPU for faster training. cfg . MODEL . DEVICE = device ### Where to save our checkpoints and logs. cfg . OUTPUT_DIR = output_dir ### Give the config back to the caller. return cfg ### main() will launch training on Ubuntu / WSL with GPU. def main (): ### Make sure the output directory exists. os . makedirs ( output_dir , exist_ok =True ) ### Build a training config with our dataset names, class count, and GPU device. cfg = get_train_cfg ( config_file_path , model_name , train_dataset_name , val_dataset_name , num_classes , device , output_dir ) ### Save the config so we can reuse it for inference later. cfg_save_path = " My-Train-DetectRon2-Ubuntu/IS_cfg-Ubunto.pickle " #IS -> Instance segmentation ### Dump the config object to disk using pickle. with open ( cfg_save_path , ' wb ' ) as f : pickle . dump ( cfg , f , protocol = pickle . HIGHEST_PROTOCOL ) # this will save our cfg ### Create a Detectron2 DefaultTrainer with our GPU config. trainer = DefaultTrainer ( cfg ) ### Load from scratch and start training. trainer . resume_or_load ( resume =False ) trainer . train () ### (Optional) you can test after training using the final checkpoint. # After training, add a testing step #trainer.test(ckpt=None) # Use the last checkpoint by default ### Standard Python pattern to run main(). if __name__ == ' __main__ ' : main () ### After training finishes, copy the output directory back to Windows ### so you can run inference there with the saved weights. # After train : #Copy the My-Train-DetectRon2 folder to the Windows enviroment to continute to the test new image Summary of this part:
We created a GPU-ready Detectron2 training workflow in Ubuntu / WSL. We trained Mask R-CNN for more iterations on cuda, which improves accuracy and speed. We saved a config pickle (IS_cfg-Ubunto.pickle) we’ll reuse for inference. The trained weights (model_final.pth) are now sitting in the output folder and ready for testing. This directly supports “mask r-cnn detectron2 tutorial” and “train detectron2 custom dataset.” I also cover how to run object detection on recorded video streams using Python and OpenCV here: https://eranfeit.net/how-to-classify-objects-in-videos-using-jetson-nano-opencv-python/
Testing the trained Detectron2 model on new images This final technical step loads the saved config, points it to the final trained weights, runs inference on new test images, and displays the predictions with colored masks and labels.
We reload the pickle file we saved during training (the config).model_final.pth) and we set a confidence threshold.
We create a DefaultPredictor, which is a simple wrapper that:
Preprocesses the image Runs the model Returns predicted classes, boxes, masks, and scores We define the class names list CLASSES.
We use Visualizer again, but now we draw predicted instance masks instead of ground truth.
Here is the full inference / testing script, annotated:
### Import DefaultPredictor to run inference with our trained config. from detectron2 . engine import DefaultPredictor ### OS and pickle for loading saved config objects and paths. import os import pickle ### Load the saved Detectron2 config from training. cfg_saved_path = " My-Train-Detectron2/IS_cfg.pickle " # IS -> Instance segmentation with open ( cfg_saved_path , ' rb ' ) as f : ### Read the config back into memory. cfg = pickle . load ( f ) # get the configuration file ### Define where our training artifacts (weights) were stored. output_dir = " My-Train-Detectron2 " ### Point the config to the final trained weights. cfg . MODEL . WEIGHTS = os . path . join ( output_dir , " model_final.pth " ) ### Set the confidence threshold for predictions (only show objects above 0.5). cfg . MODEL . ROI_HEADS . SCORE_THRESH_TEST = 0.5 # over 50% , which object should display ### Force inference to run on CPU in this script (you can set "cuda" if available). cfg . MODEL . DEVICE = " cpu " ### Build a predictor object from the config. This will run forward passes for us. predictor = DefaultPredictor ( cfg ) ### Class names in the same order Detectron2 learned them. CLASSES = [ " Apple " , " Strawberry " , " Ornage " , " Grapes " , " Banana " , " Lemon " ] ### Define paths of test images (unlabeled images). #path of the test images image_path_1 = " Train-custom-Object-Detection-model/Fruits_for_detectron2/Test/apples-vs-bananas.jpg " image_path_2 = " Train-custom-Object-Detection-model/Fruits_for_detectron2/Test/pexels-pixabay-70746.jpg " ### We'll run predictions and visualize the first test image. # show the predicitions on the test image ### Import OpenCV for reading and displaying images, NumPy for array ops, ### and Visualizer to draw model predictions. import cv2 import numpy as np from detectron2 . utils . visualizer import Visualizer ### Read the first test image. im = cv2 . imread ( image_path_1 ) ### Run the trained model on this image. outputs = predictor ( im ) ### Print the raw output dict so we can see boxes, masks, scores, and classes. print ( " ========================================= " ) print ( outputs ) print ( " ========================================= " ) ### Extract predicted class IDs as a tensor and move them to CPU. pred_classes = outputs [ ' instances ' ]. pred_classes . cpu () print ( " Pred Classes : " ) print ( pred_classes ) ### Convert the tensor to a NumPy array for easier handling. pred_classes = pred_classes . numpy () ### Flag will tell us if we detected at least one object. flag = np . size ( pred_classes ) print ( " Flag : " ) print ( flag ) ### If we found at least one object, grab the first prediction. if flag > 0 : pred_classes = pred_classes [ 0 ] # grab the first elemnet print ( " pred_classes: " ) print ( pred_classes ) ### Map the class ID (0,1,2,...) back to a human-readable fruit label. print ( CLASSES [ pred_classes ]) ### Convert the image to RGB for visualization. img_rgb = cv2 . cvtColor ( im , cv2 . COLOR_BGR2RGB ) ### Create a Visualizer for predictions instead of ground truth. v = Visualizer ( img_rgb , metadata = {}, scale = 0.6 ) # init the visualizer ### Draw the predicted instance masks and boxes. v = v . draw_instance_predictions ( outputs [ " instances " ]. to ( " cpu " )) ### Convert RGB back to BGR so OpenCV can show it. img_bgr = cv2 . cvtColor ( v . get_image (), cv2 . COLOR_RGB2BGR ) ### Display the prediction window with colored masks. cv2 . imshow ( " v " , img_bgr ) cv2 . waitKey ( 0 ) else : ### If nothing was detected above threshold, say so. print ( " Pred_classes is empty " ) ### Close any OpenCV windows after we're done. cv2 . destroyAllWindows () Summary of this part:
We reloaded the trained Detectron2 config and attached our final weights. We used DefaultPredictor to run inference and printed predicted classes. We visualized predicted instance segmentation masks for fruit objects. Now we have working Detectron2 instance segmentation end-to-end: dataset → training → prediction. FAQ : What is Detectron2 instance segmentation? Detectron2 instance segmentation predicts a pixel mask and a class label for every object, so you know not only what each object is but also its exact shape.
Why are we using Mask R-CNN? Mask R-CNN is a proven architecture for instance segmentation. It can detect each fruit and draw an accurate mask around it in the same forward pass.
Do I really need COCO format? Yes. Detectron2 expects COCO-style JSON for images, categories, boxes, and masks. Using COCO format keeps training smooth and avoids custom loaders.
Can I train Detectron2 on CPU? Yes, but it will be slower. CPU training is fine for demos and very small datasets, and it helps you debug config issues before moving to GPU.
When should I switch to GPU? Switch to GPU when you have more data, want faster iteration, or care about accuracy. GPU training with more iterations usually gives better results.
What does MAX_ITER control? MAX_ITER sets how long Detectron2 will train. A small value is good for testing. A larger value gives the model more time to learn your dataset.
Why do we save the config as a pickle file? Saving the config makes inference easy later. You reload the pickle, attach final weights, and you’re instantly ready to run predictions.
How do I map class IDs back to names? Create a CLASSES list in the same order you used for training. Then use the predicted class ID as an index to print the human-friendly label.
Can this pipeline handle other objects, not just fruit? Yes. You can point it at any object category — tools, products, lab samples, etc. All you need is consistent labeled data and a matching class list.
Why are masks better than boxes for quality control? Masks let you measure exact shape, size, and surface area. That’s important in tasks like fruit grading, defect detection, or counting overlapping objects.
Conclusion: what you can do next with Detectron2 We built a full Detectron2 instance segmentation workflow around fruit images.register_coco_instances, and visually validated that the masks line up with the fruit.
We trained Mask R-CNN in two modes:
We then reloaded the trained config, attached the final weights, and used DefaultPredictor to run inference on brand-new test images.
From here you can:
Add more classes (for example: “moldy banana” vs “fresh banana”). Train longer or tweak learning rate. Deploy the predictor into a small web app or a quality-control pipeline. Reuse the same structure for any object category, not just fruit. If you want a simple segmentation technique without deep learning, I also explain K-Means image segmentation in Python here: https://eranfeit.net/python-image-segmentation-made-easy-with-opencv-and-k-means-algorithm/
Connect : ☕ Buy me a coffee — https://ko-fi.com/eranfeit
🖥️ Email : feitgemel@gmail.com
🌐 https://eranfeit.net
🤝 Fiverr : https://www.fiverr.com/s/mB3Pbb
Enjoy,
Eran
Skip to content 
