What I need to write in the line level of torch.vison for 21 classes?

Question

In this code I found , line labels = torch.ones((records.shape[0],), dtype=torch.int64) ,that there is only one class and 0 in the case of  Faster RCNN is reserved for the Background. What would be for 21 classes?
def __getitem__(self, index: int):

file_name = self.file_names[index]
records = self.data[self.data['file_name'] == file_name]

image = np.array(Image.open(file_name), dtype=np.float32)
image /= 255.0

if self.transform:
    image = self.transform(image)  
    
if self.mode != "test":
    boxes = records[['xmin', 'ymin', 'xmax', 'ymax']].values
    
    area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
    area = torch.as_tensor(area, dtype=torch.float32)

labels = torch.ones((records.shape[0],), dtype=torch.int64)
    
    iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)
    
    target = {}

target['boxes'] = boxes
    target['labels'] = labels
    target['image_id'] = torch.tensor([index])
    target['area'] = area
    target['iscrowd'] = iscrowd 
    target['boxes'] = torch.stack(list((map(torch.tensor, target['boxes'])))).type(torch.float32)

return image, target, file_name
else:
    return image, file_name

AIFahim · Accepted Answer

I solved my problem. Here are the steps I changed :

add a 'class_name' column to the "marking" DataFrame
replace 'source' with 'class_name' in data splitting (df_folds…)
modify load_image_and_boxes() function to also return labels modify load_cutmix_image_and_boxes() function to also return labels (most work to do here)
replace labels = torch.zeros((boxes.shape[0],), dtype=torch.int64) with torch.tensor(labels, dtype=torch.int64)
in get_net() function, set config.num_classes = number of classes in your dataset

What I need to write in the line level of torch.vison for 21 classes?

One Answer

Add your own answers!

Ask a Question