Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Need help in creating dataset class for multilabel object detection in a single image #943

Open
IamSparky opened this issue Jan 18, 2021 · 0 comments

Comments

@IamSparky
Copy link

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.....

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant