ReplayCAD/get_replay_data.py at master · HULEI7/ReplayCAD · GitHub

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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import json
import os
import shutil
from tqdm import tqdm
import glob


##get_visa_generate_data

with open('data/visa_meta.json', 'r') as file:
    data = json.load(file)


train_data=data["train"]

output_path = "data/visa_generate/"
input_path="data/visa/"

for class_name in tqdm(train_data):

    class_output = output_path+class_name

    if not os.path.exists(class_output):
        os.makedirs(class_output)


    for file in train_data[class_name]:

        if file['anomaly']!=0:
            print("error")

        image_path=file['img_path']

        shutil.copy(input_path+image_path, class_output)


##get_mvtec_generate_data

output_path = "data/mvtec_generate/"
input_path= "data/mvtec"

class_names = ["bottle","cable","capsule","carpet","grid","hazelnut","leather","metal_nut","pill","screw","tile","toothbrush","transistor","wood"]

for class_name in class_names:

    source_files = glob.glob(input_path+"/"+class_name+"/train/good/*.png")
    target_path = output_path+class_name

    if not os.path.exists(target_path):
        os.makedirs(target_path)

    for file in source_files:

        shutil.copy(file,target_path)