HyperKvasir, a comprehensive multi-class image and video dataset for gastrointestinal endoscopy

import torch
import os
import glob
import natsort 
from PIL import Image
import csv

class Dataset(torch.utils.data.Dataset):
  def __init__(self, data_dir, train=True ,transform = None):
	...

  def __getitem__(self,index):
    path = self.input[index]
    input = Image.open(path)
# ['z-line', 'pylorus', 'retroflex-stomach', 'esophagitis-b-d','esophagitis-a', 'bbps-0-1', 'bbps-2-3', 'cecum', 'retroflex-rectum', 'ulcerative-colitis-grade-3', 'polyps', 'ulcerative-colitis-grade-2-3', 'ulcerative-colitis-grade-2', 'ulcerative-colitis-grade-0-1', 'ulcerative-colitis-grade-1-2', 'ulcerative-colitis-grade-1']
#     #label
    if 'z-line' in path:
        label = 0
    elif 'pylorus' in path:
        label = 1
    elif 'retroflex-stomach' in path:
        label = 2
    elif 'esophagitis' in path:
        label = 3
    elif 'bbps-0-1' in path:
        label = 4
    elif 'bbps-2-3' in path:
        label = 5
    elif 'cecum' in path:
        label = 6
    elif 'retroflex-rectum' in path:
        label = 7
    elif 'ulcerative-colitis' in path:
        label = 8
    elif 'polyps' in path:
        label = 9
    else : 
        label = None

    return data_img, label