[AI 프로젝트] 한식 분류기(5/6)

Google Drive 연동하기

from google.colab import drive
drive.mount('/gdrive', force_remount=True)

google drive를 현재 실행중인 colab에 mount하는 코드이며, 실행할 경우 어떤 google drive 계정과 연동할지 선택할 수 있다. (기존에는 인증키를 입력하는 방식이었으나 최근에 바뀐 것 같음.)

압축풀기

이런 어마어마한 크기의 압축 파일들을 압축 해제 해줘야 한다. 해당 directory 내 압축 파일을 읽어서 압축 해제하는 코드를 구현했다. 경로에 한글이 포함되어 있어서 굉장히 골치 아팠다. (총 847GB이나 되는 대용량이기 때문에 이 작업도 굉장히 오래 걸렸다.)

Dataset & Dataloader

Custom Dataset

class CustomImageDataset(Dataset):
  def __init__(self, images_path, annotations_path, code_path, transform=None, target_transform=None):
    self.img_path = pd.read_csv(images_path)
    self.annotations_path = pd.read_csv(annotations_path)
    file = open(code_path, "r")
    self.code = [file.readline().strip() for i in range(3165)]
    file.close()
    self.transform = transform
    self.target_transform = target_transform
  
  def __len__(self):
    return len(self.img_path)
  
  def __getitem__(self, idx):
    img = read_image(self.img_path.iloc[idx, 0])
    with open(self.annotations_path.iloc[idx, 0], 'r') as f:
      annotation = json.load(f)
    if self.transform:
      img = self.transform(img)
    label = self.code.index(annotation[0]["Code Name"][0:7])
    return img, label

image file path와, annotation file path에 대한 정보가 주어지면 이미지 읽어 transform 한 뒤 반환한다. transform 인자를 통해 이미지 크기 조절 및 정규화할 수 있다.

DataLoader

img_path = '/gdrive/MyDrive/proj/data/food_images.txt'
anno_path = '/gdrive/MyDrive/proj/data/food_annotations.txt'
code_path = '/gdrive/MyDrive/proj/data/Standard_Code.txt'

dataset = CustomImageDataset(images_path=img_path, 
                                  annotations_path=anno_path,
                                  code_path = code_path)

이제 dataset을 torch.utils.data.DataLoader를 이용해 load하여 학습에 사용하면 된다.