import os
import json
from PIL import Image
import torch
from torch.utils.data import Dataset

class VG200Classification(Dataset):
    def __init__(self, image_dir, phase, input_transform=None):
        self.image_dir = image_dir
        self.phase = phase
        self.ann_files = os.path.join(self.image_dir, '{}_vg200.json'.format(self.phase))
        self.load_anns()
        
        self.input_transform = input_transform
        
    def load_anns(self):
        if not os.path.exists(self.ann_files):
            self.extract_anns()
        self.anns = []
        json_data = json.load(open(self.ann_files, "r"))
        self.anns += json_data
    
    def extract_anns(self):
        raise NotImplementedError("No annotation file found at %s" % self.ann_files)
            
    def __len__(self):
        return len(self.anns)
    
    def __getitem__(self, idx):
        idx = idx % len(self)
        ann = self.anns[idx]
        img = Image.open(ann["img_path"]).convert("RGB")
        if self.input_transform:
            img = self.input_transform(img)
        return img, torch.Tensor(ann["target"])
    