Return batch with cls and bbox for custom semi-supervised training

[Bruno-Messias]

Hello Everyone!

I am trying to use YoloV8 for semi-supervised training, that is, to use data with segmentation and non-segmentation information in the same model. But, I encountered a difficulty during the customization of the trainers. I understood, after some tests, that the dataset returned depends on the task of the trainer class, and I think that this information is passed to the class get_dataloader, by the self.args, like below:

def get_dataloader(self, dataset_path, batch_size, mode="train", rank=0):
        # TODO: manage splits differently
        # calculate stride - check if model is initialized
        gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
        data =  create_dataloader(path=dataset_path,
                                 imgsz=self.args.imgsz,
                                 batch_size=batch_size,
                                 stride=gs,
                                 hyp=vars(self.args),
                                 augment=mode == "train",
                                 cache=self.args.cache,
                                 pad=0 if mode == "train" else 0.5,
                                 rect=self.args.rect,
                                 rank=rank,
                                 workers=self.args.workers,
                                 close_mosaic=self.args.close_mosaic != 0,
                                 prefix=colorstr(f'{mode}: '),
                                 shuffle=mode == "train",
                                 seed=self.args.seed)[0] if self.args.v5loader else \
            build_dataloader(self.args, batch_size, img_path=dataset_path, stride=gs, rank=rank, mode=mode)[0]

When the task is segmentation and a dataset with .yaml with segmentation, for example coco128-seg.yaml, the batchs returns a dict with bboxes, cls, mask and all the necessary information for its training.

However, when I pass a training dataset without the segmented data, using the same trainer for a segmentation task, such as coco128.yaml, the returned batch comes with empty cls and bboxes.

Could you help me how and where to change in the dataloader to return cls and bboxes, even with empty masks when I pass a dataset without segmentation data, to the same segmentation trainer?

[AyushExel]

@Laughing-q might be able to help you with this

[Laughing-q]

@Bruno-Messias @AyushExel well that's because this sorted_idx is empty when there're no segment labels.

ultralytics/ultralytics/yolo/data/augment.py

Line 653 in a58f766

masks, sorted_idx = polygons2masks_overlap((h, w), segments, downsample_ratio=self.mask_ratio)

update the following code

ultralytics/ultralytics/yolo/data/augment.py

Lines 655 to 656 in a58f766

	instances = instances[sorted_idx]
	cls = cls[sorted_idx]

to:

if len(sorted_idx):
    instances = instances[sorted_idx]
    cls = cls[sorted_idx]

should get what you want.

[Bruno-Messias]

Thank you @Laughing-q ! It worked out exactly what I wanted