TestDataset error
yangninghua opened this issue · comments
yangninghua commented
class TestDataset(Dataset):
......
def _get_image_files(self, path, ext={'.jpg', '.png'}):
images = []
for root, dirs, files in os.walk(path):
if "check" in root:
continue
print('loading image files ' + root)
# for file in files:
# if os.path.splitext(file)[-1] in ext and 'checkpoint' not in file:
# images.append(os.path.join(root, file))
for dir in dirs:
for file in os.listdir(os.path.join(path, dir)):
if os.path.splitext(file)[-1] in ext and 'checkpoint' not in file:
images.append(os.path.join(root, dir, file))
return sorted(images)
DeepKnowledge commented
it seems we have to create a black mask for each good images (in ground_truth)
GraceKafuu commented
import numpy as np
import cv2
import os
if name == 'main':
cls_name = ['carpet', 'grid', 'leather', 'tile', 'wood', 'bottle', 'cable', 'capsule',
'hazelnut', 'metal_nut', 'pill', 'screw', 'toothbrush', 'transistor', 'zipper']
path = r"F:\MVTec_DATA\data"
for d in cls_name:
subdir = path + "\\{}".format(d)
good_subdir = subdir + "\\train\\good"
savePath = subdir + "\\ground_truth\\good"
print(subdir, good_subdir, savePath)
os.makedirs(savePath, exist_ok=True)
goodList = os.listdir(good_subdir)
nimg = len(goodList)
for img in goodList:
imgName = img.split(".")[0]
imgAbsPath = good_subdir + "\\" + img
cv2img = cv2.imread(imgAbsPath)
H, W, _ = cv2img.shape
newImg = np.zeros((H, W), np.uint8)
for x in range(H):
for y in range(W):
#newImg[x, y] = [0, 0, 0]
newImg.itemset((x, y), 0)
cv2.imwrite(r"{}\{}_mask.png".format(savePath, imgName), newImg)
print("Saving-->" + r"{}\{}_mask.png".format(savePath, imgName))
I use this code to create good images mask, and can train