File size: 943 Bytes
e5c7710
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from typing import List, Tuple
from datasets import load_dataset
from PIL import Image
import os
import shutil

dataset = load_dataset("RGBD-SOD/test", "v1", split="train", cache_dir="data")
SAMPLES_DIR = "samples"


def prepare_samples():
    samples: List[Tuple[str, str, str]] = []
    for sample in dataset:
        rgb: Image.Image = sample["rgb"]
        depth: Image.Image = sample["depth"]
        gt: Image.Image = sample["gt"]
        name: str = sample["name"]
        dir_path = os.path.join(SAMPLES_DIR, name)
        shutil.rmtree(dir_path, ignore_errors=True)
        os.makedirs(dir_path, exist_ok=True)
        rgb_path = os.path.join(dir_path, f"rgb.jpg")
        rgb.save(rgb_path)
        depth_path = os.path.join(dir_path, f"depth.jpg")
        depth.save(depth_path)
        gt_path = os.path.join(dir_path, f"gt.png")
        gt.save(gt_path)

        samples.append([rgb_path, depth_path, gt_path])
    return samples