File size: 1,114 Bytes
98c1ffc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import csv
from pathlib import Path
import torchvision
def main(roco_root: str):
root = Path(roco_root)
check_images(
root / 'train/radiology', 'traindata.csv', 'train.csv'
)
check_images(
root / 'validate/radiology', 'valdata.csv', 'validate.csv'
)
check_images(
root / 'test/radiology', 'testdata.csv', 'test.csv'
)
def check_images(split_dir: Path, input_csv: str, target_output: str):
with open(split_dir / input_csv, 'r') as buf:
csv_reader = csv.reader(buf)
next(csv_reader, None)
filtered_csv = []
for row in csv_reader:
image_path = split_dir / 'images' / row[1]
try:
torchvision.io.read_image(str(image_path))
except:
continue
filtered_csv.append(row)
with open(split_dir / target_output, 'w') as csvfile:
spamwriter = csv.writer(csvfile)
for row in filtered_csv:
spamwriter.writerow(row)
if __name__ == '__main__':
main('/home/shpotes/medclip/data/roco-dataset')
{mode:full,isActive:false}
|