Dataset file structure (this corresponds to the filepath used in the code)
wheat:
|- test (contains .jpg test imgs)
|- train (contains .jpg train imgs)
|- train.csv
(all those above are from the kaggle dataset)
|- train_mask (new empty directory, will store the filtered images)