Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
hackassin
GitHub Repository: hackassin/learnopencv
 Path: blob/master/Bag-Of-Tricks-For-Image-Classification/split_food-101.py
3119 views
1
import argparse

2
import os

3
import os.path as osp

4
from shutil import copyfile

5


6
from tqdm import tqdm

7


8


9
def main():

10
    parser = argparse.ArgumentParser(

11
        description="Separate Food-101 into train/test folders",

12
    )

13
    parser.add_argument(

14
        "--data-root",

15
        default="./data",

16
        type=str,

17
        help="Path to root folder of the dataset",

18
    )

19
    args = parser.parse_args()

20
    classes = [

21
        "apple_pie",

22
        "bruschetta",

23
        "caesar_salad",

24
        "steak",

25
        "spring_rolls",

26
        "spaghetti_carbonara",

27
        "frozen_yogurt",

28
        "falafel",

29
        "mussels",

30
        "ramen",

31
        "onion_rings",

32
        "oysters",

33
        "risotto",

34
        "waffles",

35
        "cup_cakes",

36
        "grilled_cheese_sandwich",

37
        "fried_calamari",

38
        "huevos_rancheros",

39
        "croque_madame",

40
        "bread_pudding",

41
        "dumplings",

42
    ]

43
    assert osp.isdir(args.data_root)

44
    assert "images" in os.listdir(args.data_root)

45
    assert "meta" in os.listdir(args.data_root)

46
    os.makedirs(osp.join(args.data_root, "train"), exist_ok=True)

47
    os.makedirs(osp.join(args.data_root, "test"), exist_ok=True)

48
    for cls_name in classes:

49
        os.makedirs(osp.join(args.data_root, "train", cls_name), exist_ok=True)

50
        os.makedirs(osp.join(args.data_root, "test", cls_name), exist_ok=True)

51
    with open(osp.join(args.data_root, "meta", "train.txt"), "r") as file:

52
        for image in tqdm(file):

53
            image = image.rstrip()

54
            if image.split("/")[0] in classes:

55
                copyfile(

56
                    osp.join(args.data_root, "images", image + ".jpg"),

57
                    osp.join(args.data_root, "train", image + ".jpg"),

58
                )

59
    with open(osp.join(args.data_root, "meta", "test.txt"), "r") as file:

60
        for image in tqdm(file):

61
            image = image.rstrip()

62
            if image.split("/")[0] in classes:

63
                copyfile(

64
                    osp.join(args.data_root, "images", image + ".jpg"),

65
                    osp.join(args.data_root, "test", image + ".jpg"),

66
                )

67


68


69
if __name__ == "__main__":

70
    main()

71


72