YoungGod / DFR

Project: Unsupervised Anomaly Segmentation via Deep Feature Reconstruction

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

TestDataset error

yangninghua opened this issue · comments

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)

it seems we have to create a black mask for each good images (in ground_truth)

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