sam-tp / quickstart.py
jamiewjm's picture
Upload quickstart.py
e1f9f79 verified
from datasets import load_dataset
from pathlib import Path
from PIL import Image
REPO = "your_user/sam-tp" # change to your dataset repo id
ds_imgs = load_dataset(
"imagefolder",
data_dir=".",
data_files={"image": f"hf://datasets/{REPO}/images/**"},
split="train",
)
ds_msks = load_dataset(
"imagefolder",
data_dir=".",
data_files={"mask": f"hf://datasets/{REPO}/annotations/**"},
split="train",
)
mask_index = {Path(r["image"]["path"]).name: r["image"]["path"] for r in ds_msks}
row = ds_imgs[0]
img_path = Path(row["image"]["path"])
msk_path = Path(mask_index[img_path.name])
print("Image:", img_path)
print("Mask: ", msk_path)
Image.open(img_path).show()
Image.open(msk_path).show()