open-images-dataset.py

class OpenImages(Dataset):
  def __init__(self, image_folder, df):
    self.root = image_folder
    self.df = df
    self.unique_images = df['ImageID'].unique()
  def __len__(self): return len(self.unique_images)
  def __getitem__(self, ix):
    image_id = self.unique_images[ix]
    image_path = f'{self.root}/{image_id}.jpg'
    image = cv2.imread(image_path, 1)[...,::-1] # conver BGR to RGB
    h, w, _ = image.shape
    df = self.df.copy()
    df = df[df['ImageID'] == image_id]
    boxes = df['XMin,YMin,XMax,YMax'.split(',')].values
    boxes = (boxes * np.array([w,h,w,h])).astype(np.uint16).tolist()
    classes = df['LabelName'].values.tolist()
    return image, boxes, classes