-
Notifications
You must be signed in to change notification settings - Fork 4.1k
Open
Description
Can some please help me in modifying this single class object detection dataset class to multiple class object detection dataset with bounding boxes annotations for 15 classes , with 0 being the first class and 14th class will be the image with none of the classes given.
import os
import numpy as np
import torch
from PIL import Image
class object_detection_dataset_class(Dataset):
def __init__(self, csv_file_path):
self.train_df = pd.read_csv(csv_file_path)
self.image_ids = self.train_df['image_id'].unique()
def __len__(self):
return len(self.image_ids.shape[0])
def __getitem__(self, index):
image_id = self.image_ids[index]
bboxes = self.train_df[self.train_df['image_id'] == image_id]
image = np.array(Image.open('__image_file_path__'+ self.image_id[index] +'.jpg'))
image = Image.fromarray(image).convert('RGB')
image /= 255.0
boxes = bboxes[['xmin', 'ymin', 'xmax', 'ymax']]
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
boxes = torch.as_tensor(boxes, dtype = torch.float32)
area = torch.as_tensor(area, dtype = torch.float32)
labels = torch.ones((bboxes.shape[0], ), dtype = torch.int64)
is_crowd = torch.ones((bboxes.shape[0], ), dtype = torch.int64)
target = {}
target["boxes"] = boxes
target["labels"] = labels
target["masks"] = masks
target["image_id"] = image_id
target["area"] = area
target["iscrowd"] = iscrowd
image = torchvision.transforms.ToTensor()(image)
return image, target
initializing the dataset
train_dataset = object_detection_dataset_class("______csv file path for training data___________")
def collate_fn(batch):
return tuple(zip(*batch)
training_data_loader = Dataloader(train_dataset, batch_size = 2, shuffle= True, num_workers= 2, collate_fn = collate_fn)
for dry run testing of dataloader to see the images loaded inside of pytorch dataloader
from matplotlib import pyplot as plt
images, targets = next(iter(training_data_loader))
images = list(image.to(device) for image in images)
targets = [{k: v.to(device) for k,v in t.items()} for t in targets]
boxes = targets[0]['boxes'].cpu().numpy().astype(np.int32)
img = images[0].permute(1, 2, 0).cpu().numpy()
fig, ax = plt.subplots(1,1, figsize(12, 6))
for box in boxes:
cv2.rectangle(image,
(box[0], box[1]),
(box[2], box[3])
(220, 0, 0),
1)
ax.set_axis_off()
ax.imshow(img)
Early thanks if anyone can help me on this.....
Metadata
Metadata
Assignees
Labels
No labels