# Copyright 2022 Cristóbal Alcázar
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""NIH Chest X-ray Dataset"""
import os
import datasets
from requests import get
from pandas import read_csv
logger = datasets.logging.get_logger(__name__)
_CITATION = """\
@inproceedings{Wang_2017,
doi = {10.1109/cvpr.2017.369},
url = {https://doi.org/10.1109%2Fcvpr.2017.369},
year = 2017,
month = {jul},
publisher = {{IEEE}
},
author = {Xiaosong Wang and Yifan Peng and Le Lu and Zhiyong Lu and Mohammadhadi Bagheri and Ronald M. Summers},
title = {{ChestX}-Ray8: Hospital-Scale Chest X-Ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases},
booktitle = {2017 {IEEE} Conference on Computer Vision and Pattern Recognition ({CVPR})}
}
"""
_DESCRIPTION = """\
The NIH Chest X-ray dataset consists of 100,000 de-identified images of chest x-rays. The images are in PNG format.
The data is provided by the NIH Clinical Center and is available through the NIH download site: https://nihcc.app.box.com/v/ChestXray-NIHCC
"""
_HOMEPAGE = "https://nihcc.app.box.com/v/chestxray-nihcc"
_REPO = "https://huggingface.co/datasets/alkzar90/NIH-Chest-X-ray-dataset/resolve/main/data"
_IMAGE_URLS = [
f"{_REPO}/images/images_001.zip",
f"{_REPO}/images/images_002.zip",
f"{_REPO}/images/images_003.zip",
f"{_REPO}/images/images_004.zip",
f"{_REPO}/images/images_005.zip",
f"{_REPO}/images/images_006.zip",
f"{_REPO}/images/images_007.zip",
f"{_REPO}/images/images_008.zip",
f"{_REPO}/images/images_009.zip",
f"{_REPO}/images/images_010.zip",
f"{_REPO}/images/images_011.zip",
f"{_REPO}/images/images_012.zip"
#'https://huggingface.co/datasets/alkzar90/NIH-Chest-X-ray-dataset/resolve/main/dummy/0.0.0/images_001.tar.gz',
#'https://huggingface.co/datasets/alkzar90/NIH-Chest-X-ray-dataset/resolve/main/dummy/0.0.0/images_002.tar.gz'
]
_URLS = {
"train_val_list": f"{_REPO}/train_val_list.txt",
"test_list": f"{_REPO}/test_list.txt",
"labels": f"{_REPO}/Data_Entry_2017_v2020.csv",
"BBox": f"{_REPO}/BBox_List_2017.csv",
"image_urls": _IMAGE_URLS
}
_LABEL2IDX = {"No Finding": 0,
"Atelectasis": 1,
"Cardiomegaly": 2,
"Effusion": 3,
"Infiltration": 4,
"Mass": 5,
"Nodule": 6,
"Pneumonia": 7,
"Pneumothorax": 8,
"Consolidation": 9,
"Edema": 10,
"Emphysema": 11,
"Fibrosis": 12,
"Pleural_Thickening": 13,
"Hernia": 14}
_NAMES = list(_LABEL2IDX.keys())
class ChestXray14Config(datasets.BuilderConfig):
"""NIH Image Chest X-ray14 configuration."""
def __init__(self, name, **kwargs):
super(ChestXray14Config, self).__init__(
version=datasets.Version("1.0.0"),
name=name,
description="NIH ChestX-ray14",
**kwargs,
)
class ChestXray14(datasets.GeneratorBasedBuilder):
"""NIH Image Chest X-ray14 dataset."""
BUILDER_CONFIGS = [
ChestXray14Config("image-classification"),
ChestXray14Config("object-detection"),
]
def _info(self):
if self.config.name == "image-classification":
features = datasets.Features(
{
"image": datasets.Image(),
"labels": datasets.features.Sequence(
datasets.features.ClassLabel(
num_classes=len(_NAMES),
names=_NAMES
)
),
}
)
keys = ("image", "labels")
if self.config.name == "object-detection":
features = datasets.Features(
{
"image_id": datasets.Value("string"),
"patient_id": datasets.Value("int32"),
"image": datasets.Image(),
"width": datasets.Value("int32"),
"height": datasets.Value("int32"),
}
)
object_dict = {
"image_id": datasets.Value("string"),
"area": datasets.Value("int64"),
"bbox": datasets.Sequence(datasets.Value("float32"), length=4),
}
features["objects"] = [object_dict]
keys = ("image", "objects")
return datasets.DatasetInfo(
description=_DESCRIPTION,
features=features,
supervised_keys=keys,
homepage=_HOMEPAGE,
citation=_CITATION,
)
def _split_generators(self, dl_manager):
# Get the image names that belong to the train-val dataset
logger.info("Downloading the train_val_list image names")
train_val_list = get(_URLS['train_val_list']).iter_lines()
train_val_list = set([x.decode('UTF8') for x in train_val_list])
logger.info(f"Check train_val_list: {train_val_list}")
# Create list for store the name of the images for each dataset
train_files = []
test_files = []
# Download batches
data_files = dl_manager.download_and_extract(_URLS["image_urls"])
# Iterate trought image folder and check if they belong to
# the trainset or testset
for batch in data_files:
logger.info(f"Batch for data_files: {batch}")
path_files = dl_manager.iter_files(batch)
for img in path_files:
if os.path.basename(img) in train_val_list:
train_files.append(img)
else:
test_files.append(img)
return [
datasets.SplitGenerator(
name=datasets.Split.TRAIN,
gen_kwargs={
"files": train_files
}
),
datasets.SplitGenerator(
name=datasets.Split.TEST,
gen_kwargs={
"files": test_files
}
)
]
def _generate_examples(self, files):
if self.config.name == "image-classification":
# Read csv with image labels
label_csv = read_csv(_URLS["labels"])
for i, path in enumerate(files):
file_name = os.path.basename(path)
# Get image id to filter the respective row of the csv
image_id = file_name
image_labels = label_csv[label_csv["Image Index"] == image_id]["Finding Labels"].values[0].split("|")
if file_name.endswith(".png"):
yield i, {
"image": path,
"labels": image_labels,
}
