class CreateSegmentionSample():
def __init__(self):
self.dirProcess = DirProcess()
self.image_process = ImageProcess()
self.annotation_post = ".png"
def create_train_and_test(self, inputDir, outputPath, probability):
annotationsDir = os.path.join(inputDir, "../SegmentLabel")
saveTrainFilePath = os.path.join(outputPath, "train.txt")
saveTestFilePath = os.path.join(outputPath, "val.txt")
saveTrainFilePath = open(saveTrainFilePath, "w")
saveTestFilePath = open(saveTestFilePath, "w")
imageList = list(self.dirProcess.getDirFiles(inputDir, "*.*"))
random.shuffle(imageList)
for imageIndex, imagePath in enumerate(imageList):
print(imagePath)
image = cv2.imdecode(np.fromfile(imagePath, dtype=np.uint8),
cv2.IMREAD_GRAYSCALE)
path, file_name_and_post = os.path.split(imagePath)
imageName, post = os.path.splitext(file_name_and_post)
seg_label_name = imageName + self.annotation_post
label_path = os.path.join(annotationsDir, seg_label_name)
if (image is not None) and os.path.exists(label_path):
if (imageIndex + 1) % probability == 0:
saveTestFilePath.write("%s\n" % file_name_and_post)
else:
saveTrainFilePath.write("%s\n" % file_name_and_post)
saveTrainFilePath.close()
saveTestFilePath.close()
class ConvertSegmentionLable():
def __init__(self):
self.save_label_dir = "SegmentLabel"
self.annotation_post = ".png"
self.dirProcess = DirProcess()
self.image_process = ImageProcess()
def convert_segment_label(self, label_dir, is_gray, class_list):
output_dir = os.path.join(label_dir, "../%s" % self.save_label_dir)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
label_list = list(self.dirProcess.getDirFiles(label_dir, "*.*"))
for label_path in label_list:
path, file_name_and_post = os.path.split(label_path)
print(label_path)
mask = self.process_segment_label(label_path, is_gray, class_list)
if mask is not None:
save_path = os.path.join(output_dir, file_name_and_post)
cv2.imwrite(save_path, mask)
def process_segment_label(self, label_path, is_gray, class_list):
if is_gray:
mask = self.image_process.read_gray_image(label_path)
else:
_, mask = self.image_process.readRgbImage(label_path)
if mask is not None:
if is_gray:
mask = self.convert_gray_label(mask, class_list)
else:
mask = self.convert_color_label(mask, class_list)
return mask
def convert_gray_label(self, mask, class_list):
shape = mask.shape # shape = [height, width]
result = np.full(shape, 250, dtype=np.uint8)
for index, value in enumerate(class_list):
gray_value = int(value[1].strip())
result[mask == gray_value] = index
return result
def convert_color_label(self, mask, class_list):
shape = mask.shape[:2] # shape = [height, width]
result = np.full(shape, 250, dtype=np.uint8)
for index, value in enumerate(class_list):
value_list = [int(x) for x in value[1].spilt(',') if x.strip()]
color_value = np.array(value_list, dtype=np.uint8)
temp1 = mask[:, :] == color_value
temp2 = np.sum(temp1, axis=2)
result[temp2 == 3] = index
return result
class ImagesLoader(DataLoader):
def __init__(self, input_dir, image_size=(416, 416)):
super().__init__()
self.image_size = image_size
self.imageProcess = ImageProcess()
self.dirProcess = DirProcess()
self.dataset_process = ImageDataSetProcess()
temp_files = self.dirProcess.getDirFiles(input_dir, "*.*")
self.files = list(temp_files)
self.count = len(self.files)
self.color = (127.5, 127.5, 127.5)
def __iter__(self):
self.index = -1
return self
def __next__(self):
self.index += 1
if self.index == self.count:
raise StopIteration
image_path = self.files[self.index]
# Read image
srcImage, rgb_image = self.imageProcess.readRgbImage(image_path)
# Padded resize
rgb_image, _, _ = self.dataset_process.image_resize_square(rgb_image,
self.image_size,
self.color)
rgb_image = self.dataset_process.image_normaliza(rgb_image)
numpy_image = self.dataset_process.numpy_transpose(rgb_image)
torch_image = self.all_numpy_to_tensor(numpy_image)
return srcImage, torch_image
def __len__(self):
return self.count
class CreateDetectionSample():
def __init__(self):
self.dirProcess = DirProcess()
self.xmlProcess = XMLProcess()
self.annotation_post = ".xml"
def createBalanceSample(self, inputTrainPath, outputPath):
if not os.path.exists(outputPath):
os.makedirs(outputPath)
path, _ = os.path.split(inputTrainPath)
annotationDir = os.path.join(path, "../Annotations")
imagesDir = os.path.join(path, "../JPEGImages")
writeFile = self.createWriteFile(outputPath)
for fileNameAndPost in self.dirProcess.getFileData(inputTrainPath):
fileName, post = os.path.splitext(fileNameAndPost)
annotationFileName = fileName + self.annotation_post
annotationPath = os.path.join(annotationDir, annotationFileName)
imagePath = os.path.join(imagesDir, fileNameAndPost)
print(imagePath, annotationPath)
if os.path.exists(annotationPath) and \
os.path.exists(imagePath):
_, _, boxes = self.xmlProcess.parseRectData(annotationPath)
allNames = [
box.name for box in boxes
if box.name in detect2d_config.className
]
names = set(allNames)
print(names)
for className in names:
writeFile[className].write(fileNameAndPost + "\n")
def createTrainAndTest(self, inputDir, outputPath, probability):
annotationsDir = os.path.join(inputDir, "../Annotations")
saveTrainFilePath = os.path.join(outputPath, "train.txt")
saveTestFilePath = os.path.join(outputPath, "val.txt")
saveTrainFilePath = open(saveTrainFilePath, "w")
saveTestFilePath = open(saveTestFilePath, "w")
imageList = list(self.dirProcess.getDirFiles(inputDir, "*.*"))
random.shuffle(imageList)
for imageIndex, imagePath in enumerate(imageList):
print(imagePath)
image = cv2.imdecode(np.fromfile(imagePath, dtype=np.uint8),
cv2.IMREAD_GRAYSCALE)
path, file_name_and_post = os.path.split(imagePath)
imageName, post = os.path.splitext(file_name_and_post)
xmlPath = os.path.join(annotationsDir,
"%s%s" % (imageName, self.annotation_post))
if (image is not None) and os.path.exists(xmlPath):
if (imageIndex + 1) % probability == 0:
saveTestFilePath.write("%s\n" % file_name_and_post)
else:
saveTrainFilePath.write("%s\n" % file_name_and_post)
saveTrainFilePath.close()
saveTestFilePath.close()
def createWriteFile(self, outputPath):
result = {}
for className in detect2d_config.className:
classImagePath = os.path.join(outputPath, className + ".txt")
result[className] = open(classImagePath, "w")
return result