import json
import logging
import os
import torchvision.transforms as transforms
import torch.nn.functional as F
from torch.utils.data import Dataset
import numpy as np
from PIL import Image
from typing import Tuple
from datasets.utils import set_default_from_args
from datasets.utils.hf_download import ensure_required_files_from_hf, download_dataset_file
from utils import smart_joint
from utils.conf import base_path
from datasets.utils.continual_dataset import ContinualDataset, fix_class_names_order, store_masked_loaders
from datasets.transforms.denormalization import DeNormalize
from torchvision.transforms.functional import InterpolationMode
from utils.prompt_templates import templates
[docs]
class CropDisease(Dataset):
HF_REPO_ID = 'aimagelab-ta/cropdisease'
HF_REVISION = 'main'
REQUIRED_SPLITS = ['train.json', 'test.json']
PARQUET_FILES = {
'train': 'cropdisease_train.parquet',
'test': 'cropdisease_test.parquet',
}
READY_FILE = 'DONE'
LABELS = [
"Apple___Apple_scab",
"Apple___Black_rot",
"Apple___healthy",
"Blueberry___healthy",
"Cherry___Powdery_mildew",
"Cherry___healthy",
"Corn___Cercospora_leaf_spot Gray_leaf_spot",
"Corn___Common_rust",
"Corn___Northern_Leaf_Blight",
"Corn___healthy",
"Grape___Black_rot",
"Grape___Esca_(Black_Measles)",
"Grape___Leaf_blight_(Isariopsis_Leaf_Spot)",
"Grape___healthy",
"Orange___Haunglongbing_(Citrus_greening)",
"Peach___Bacterial_spot",
"Pepper,_bell___Bacterial_spot",
"Pepper,_bell___healthy",
"Potato___Early_blight",
"Potato___Late_blight",
"Raspberry___healthy",
"Soybean___healthy",
"Squash___Powdery_mildew",
"Strawberry___Leaf_scorch",
"Strawberry___healthy",
"Tomato___Bacterial_spot",
"Tomato___Early_blight",
"Tomato___Late_blight",
"Tomato___Leaf_Mold",
"Tomato___Septoria_leaf_spot",
"Tomato___Spider_mites Two-spotted_spider_mite",
"Tomato___Target_Spot",
"Tomato___Tomato_Yellow_Leaf_Curl_Virus",
"Tomato___Tomato_mosaic_virus",
"Tomato___healthy",
]
def __init__(self, root, train=True, transform=None,
target_transform=None, download=False) -> None:
self.root = root
self.train = train
self.transform = transform
self.target_transform = target_transform
self.not_aug_transform = transforms.Compose([
transforms.Resize((224, 224), interpolation=InterpolationMode.BICUBIC),
transforms.ToTensor()]
)
filename = smart_joint(root, ('train' if train else 'test') + '.json')
self._ensure_local_data(root=root, filename=filename, download=download)
with open(filename) as f:
data_config = json.load(f)
self.data = np.array([smart_joint(root, 'images', d) for d in data_config['data']])
self.targets = np.array(data_config['labels']).astype(np.int16)
@classmethod
def _download_from_legacy_source(cls, root: str) -> None:
from onedrivedownloader import download
ln = "https://unimore365-my.sharepoint.com/:u:/g/personal/215580_unimore_it/EZUaXKQUAVBPrhjHTUdflDEBNu0YiPWrdpAdDhnEU4nD2A?e=GPrCYF"
logging.info('Downloading CropDisease dataset from OneDrive')
parent_dir = os.path.dirname(root)
download(ln, filename=os.path.join(root, 'cropdisease.tar.gz'), unzip=True, unzip_path=parent_dir, clean=True)
@classmethod
def _extract_images_from_parquet(cls, root: str) -> None:
import pyarrow.parquet as pq
from tqdm.auto import tqdm
image_root = smart_joint(root, 'images')
os.makedirs(image_root, exist_ok=True)
written = 0
for parquet_name in cls.PARQUET_FILES.values():
parquet_path = smart_joint(root, parquet_name)
if not os.path.isfile(parquet_path):
raise FileNotFoundError(f'Parquet file not found: {parquet_path}')
parquet_file = pq.ParquetFile(parquet_path)
total_rows = parquet_file.metadata.num_rows if parquet_file.metadata is not None else None
pbar = tqdm(total=total_rows, desc=f'Extracting {parquet_name}', leave=False)
for batch in parquet_file.iter_batches(columns=['filename', 'image_bytes'], batch_size=512):
data = batch.to_pydict()
for relname, image_bytes in zip(data['filename'], data['image_bytes']):
out_path = smart_joint(image_root, relname)
if os.path.isfile(out_path):
continue
with open(out_path, 'wb') as f:
f.write(bytes(image_bytes))
written += 1
pbar.update(batch.num_rows)
pbar.close()
logging.info('Extracted %d CropDisease images from parquet', written)
@classmethod
def _ensure_local_data(cls, root: str, filename: str, download: bool) -> None:
os.makedirs(root, exist_ok=True)
ready_path = smart_joint(root, cls.READY_FILE)
split_paths = [smart_joint(root, split_file) for split_file in cls.REQUIRED_SPLITS]
if os.path.isfile(ready_path) and all(os.path.isfile(path) for path in split_paths):
return
if not download:
raise FileNotFoundError(
f'Missing CropDisease metadata in `{root}`. '
f'Expected `{cls.READY_FILE}` and split files {cls.REQUIRED_SPLITS}.'
)
try:
missing_splits = [
split_name for split_name, split_path in zip(cls.REQUIRED_SPLITS, split_paths)
if not os.path.isfile(split_path)
]
if missing_splits:
ensure_required_files_from_hf(
local_dir=root,
required_relpaths=cls.REQUIRED_SPLITS,
repo_id=cls.HF_REPO_ID,
revision=cls.HF_REVISION,
)
for parquet_name in cls.PARQUET_FILES.values():
parquet_path = smart_joint(root, parquet_name)
if not os.path.isfile(parquet_path):
download_dataset_file(
repo_id=cls.HF_REPO_ID,
local_dir=root,
filename=parquet_name,
revision=cls.HF_REVISION,
)
cls._extract_images_from_parquet(root)
with open(ready_path, 'w') as f:
f.write('')
except Exception as e:
logging.warning('HF parquet download for CropDisease failed, falling back to OneDrive: %s', e)
cls._download_from_legacy_source(root)
with open(ready_path, 'w') as f:
f.write('')
if not os.path.isfile(filename):
raise FileNotFoundError(f'File not found after download: {filename}')
if not os.path.isfile(ready_path) or not all(os.path.isfile(path) for path in split_paths):
raise FileNotFoundError(
f'CropDisease dataset is not marked as ready in `{root}`. '
f'Missing `{cls.READY_FILE}` or split files {cls.REQUIRED_SPLITS}.'
)
def __len__(self):
return len(self.targets)
def __getitem__(self, index: int) -> Tuple[Image.Image, int, Image.Image]:
"""
Gets the requested element from the dataset.
:param index: index of the element to be returned
:returns: tuple: (image, target) where target is index of the target class.
"""
img, target = self.data[index], self.targets[index]
img = Image.open(img).convert('RGB')
original_img = img.copy()
not_aug_img = self.not_aug_transform(original_img)
if self.transform is not None:
img = self.transform(img)
if self.target_transform is not None:
target = self.target_transform(target)
if not self.train:
return img, target
if hasattr(self, 'logits'):
return img, target, not_aug_img, self.logits[index]
return img, target, not_aug_img
[docs]
class SequentialCropDisease(ContinualDataset):
NAME = 'seq-cropdisease'
SETTING = 'class-il'
N_TASKS = 7
N_CLASSES = 35
N_CLASSES_PER_TASK = N_CLASSES // N_TASKS
SIZE = (224, 224)
MEAN, STD = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
TRANSFORM = transforms.Compose([
transforms.RandomResizedCrop(SIZE, interpolation=InterpolationMode.BICUBIC),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=MEAN, std=STD),
])
TEST_TRANSFORM = transforms.Compose([
transforms.Resize(size=SIZE, interpolation=InterpolationMode.BICUBIC),
transforms.CenterCrop(SIZE),
transforms.ToTensor(),
transforms.Normalize(mean=MEAN, std=STD),
])
[docs]
def get_data_loaders(self):
train_dataset = CropDisease(base_path() + 'cropdisease', train=True,
download=True, transform=self.TRANSFORM)
test_dataset = CropDisease(base_path() + 'cropdisease', train=False,
download=True, transform=self.TEST_TRANSFORM)
train, test = store_masked_loaders(train_dataset, test_dataset, self)
return train, test
[docs]
def get_class_names(self):
if self.class_names is not None:
return self.class_names
classes = [x.replace('_', ' ') for x in CropDisease.LABELS] # .split('___')[-1]
classes = fix_class_names_order(classes, self.args)
self.class_names = classes
return self.class_names
[docs]
@staticmethod
def get_prompt_templates():
return templates['cifar100']
[docs]
@set_default_from_args("backbone")
def get_backbone():
return "vit"
[docs]
@staticmethod
def get_loss():
return F.cross_entropy
[docs]
@set_default_from_args('n_epochs')
def get_epochs():
return 5
[docs]
@set_default_from_args('batch_size')
def get_batch_size():
return 128