File size: 2,078 Bytes
be2715b |
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 64 65 66 67 68 69 70 71 72 73 |
import numpy as np
import torch
import shutil
import os
import matplotlib.pyplot as plt
import cv2
import json
from PIL import Image
import pickle
from skimage.transform import resize
from utils.dataset_prepare import split_data, save_fileLabel
def BUID_split():
dataset_name = "BUID"
file = "./dataset_demo/BUID"
if os.path.exists(os.path.join(file, "normal")):
shutil.rmtree(os.path.join(file, "normal"))
dir_data = "./dataset_demo/BUID/images"
dir_label = "./dataset_demo/BUID/labels"
if os.path.exists(dir_data):
shutil.rmtree(dir_data)
os.mkdir(dir_data)
if os.path.exists(dir_label):
shutil.rmtree(dir_label)
os.mkdir(dir_label)
for i in os.listdir(file):
if i == "labels" or i == "images" or "label" in i:
continue
file_label = os.path.join(file, i)
for img in os.listdir(file_label):
img_file = os.path.join(file_label, img)
if "mask" in img:
shutil.copy(img_file, os.path.join(dir_label, img))
else:
shutil.copy(img_file, os.path.join(dir_data, img))
file = os.listdir(dir_label)
label_uni = -1
check = False
b = 0
for i in os.listdir(dir_data):
for k in range(10):
if k == 0:
mask = "_mask"
else:
mask = "_mask_" + str(k)
a = i.replace(".png", mask+".png")
if a in file:
b = k
if not check:
label_uni = cv2.imread(os.path.join(dir_label, a))
check = True
else:
img = cv2.imread(os.path.join(dir_label, a))
label_uni = label_uni + img
os.remove(os.path.join(dir_label, a))
else:
check = False
break
#print(i)
cv2.imwrite(os.path.join(dir_label, i), label_uni)
label_uni = -1
check = False
b = 0
split_data(dataset_name)
save_fileLabel(dataset_name) |