537 lines
19 KiB
Python
537 lines
19 KiB
Python
import colorsys
|
||
import copy
|
||
import json
|
||
import math
|
||
import os
|
||
from pathlib import Path
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image, ImageDraw
|
||
from tabulate import tabulate
|
||
from torchvision.transforms import transforms, InterpolationMode
|
||
|
||
"""
|
||
生成num种颜色
|
||
返回值: color list，返回的color list的第一个数值永远是(0, 0, 0)
|
||
"""
|
||
def get_colors(num: int):
|
||
assert num >= 1
|
||
if num <= 21:
|
||
colors = [
|
||
(0, 0, 0),
|
||
(128, 0, 0),
|
||
(0, 128, 0),
|
||
(128, 128, 0),
|
||
(0, 0, 128),
|
||
(128, 0, 128),
|
||
(0, 128, 128),
|
||
(128, 128, 128),
|
||
(64, 0, 0),
|
||
(192, 0, 0),
|
||
(64, 128, 0),
|
||
(192, 128, 0),
|
||
(64, 0, 128),
|
||
(192, 0, 128),
|
||
(64, 128, 128),
|
||
(192, 128, 128),
|
||
(0, 64, 0),
|
||
(128, 64, 0),
|
||
(0, 192, 0),
|
||
(128, 192, 0),
|
||
(0, 64, 128),
|
||
(128, 64, 12)
|
||
]
|
||
else:
|
||
hsv_tuples = [(x / num, 1., 1.) for x in range(0, num - 1)]
|
||
colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
|
||
colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))
|
||
if (0, 0, 0) in colors:
|
||
colors.remove((0, 0, 0))
|
||
colors = [(0, 0, 0), *colors]
|
||
return colors
|
||
|
||
"""
|
||
获取某种颜色对应的标签
|
||
返回值：标签值
|
||
"""
|
||
def get_label_of_color(colors, color):
|
||
low_label = colors.index(color)
|
||
return low_label, 255 - low_label
|
||
|
||
"""
|
||
获取某个标签值对应的颜色
|
||
返回值：元组(r, g, b)
|
||
"""
|
||
def get_color_of_label(colors, label):
|
||
low_label = label if label < 255 - label else 255 - label
|
||
return colors[low_label]
|
||
|
||
"""
|
||
获取某种类别对应的标签
|
||
返回值：标签值
|
||
"""
|
||
def get_label_of_cls(classes, cls):
|
||
low_label = classes.index(cls)
|
||
return low_label, 255 - low_label
|
||
|
||
"""
|
||
获取某个标签值对应的类别
|
||
返回值：类别
|
||
"""
|
||
def get_cls_of_label(classes, label):
|
||
low_label = label if label < 255 - label else 255 - label
|
||
return classes[low_label]
|
||
|
||
"""
|
||
获取某种颜色对应的类别
|
||
返回值：类别
|
||
color: (r, g, b)
|
||
"""
|
||
def get_cls_of_color(classes, colors, color):
|
||
idx = colors.index(color)
|
||
return get_cls_of_label(classes, idx)
|
||
|
||
"""
|
||
获取某种类别对应的颜色
|
||
返回值：颜色，(r, g, b)
|
||
"""
|
||
def get_color_of_cls(classes, colors, cls):
|
||
idx = classes.index(cls)
|
||
return get_color_of_label(colors, idx)
|
||
|
||
|
||
def draw_mask(draw, points, shape_type, label, out_line_value, line_width=10, point_width=5):
|
||
points = [tuple(point) for point in points]
|
||
if shape_type == 'circle':
|
||
assert len(points) == 2, 'Shape of shape_type=circle must have 2 points'
|
||
(cx, cy), (px, py) = points
|
||
d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2)
|
||
draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=out_line_value, fill=label)
|
||
elif shape_type == 'rectangle':
|
||
assert len(points) == 2, 'Shape of shape_type=rectangle must have 2 points'
|
||
draw.rectangle(points, outline=out_line_value, fill=label)
|
||
elif shape_type == 'line':
|
||
assert len(points) == 2, 'Shape of shape_type=line must have 2 points'
|
||
greater_label = out_line_value
|
||
draw.line(xy=points, fill=greater_label, width=line_width)
|
||
elif shape_type == 'linestrip':
|
||
greater_label = out_line_value
|
||
draw.line(xy=points, fill=greater_label, width=line_width)
|
||
elif shape_type == 'point':
|
||
assert len(points) == 1, 'Shape of shape_type=point must have 1 points'
|
||
cx, cy = points[0]
|
||
r = point_width
|
||
draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=out_line_value, fill=label)
|
||
else:
|
||
assert len(points) > 2, 'Polygon must have points more than 2'
|
||
draw.polygon(xy=points, outline=out_line_value, fill=label)
|
||
|
||
"""
|
||
负责将labelme的标记转换成(mask)图像
|
||
classes: 类别列表
|
||
"""
|
||
def labelme_json2mask(classes, json_path: str, mask_saved_path: str):
|
||
assert classes is not None and classes[0] == "background"
|
||
|
||
json_path = Path(json_path)
|
||
if json_path.exists() and json_path.is_file():
|
||
with json_path.open(mode="r") as f:
|
||
json_data = json.load(f)
|
||
image_height = json_data["imageHeight"]
|
||
image_width = json_data["imageWidth"]
|
||
image_path = json_data["imagePath"]
|
||
shapes = json_data["shapes"]
|
||
|
||
cls_info_list = []
|
||
for shape in shapes:
|
||
cls_name_in_json = shape["label"]
|
||
assert cls_name_in_json in classes
|
||
points = shape["points"]
|
||
shape_type = shape["shape_type"]
|
||
label_of_cls = classes.index(cls_name_in_json)
|
||
cls_info_list.append(
|
||
{
|
||
"cls_name": cls_name_in_json,
|
||
"label": label_of_cls,
|
||
"points": points,
|
||
"shape_type": shape_type
|
||
}
|
||
)
|
||
|
||
mask = np.zeros(shape=(image_height, image_width), dtype=np.uint8)
|
||
mask = Image.fromarray(mask)
|
||
draw = ImageDraw.Draw(mask)
|
||
for cls_info in cls_info_list:
|
||
points = cls_info["points"]
|
||
shape_type = cls_info["shape_type"]
|
||
label = cls_info["label"]
|
||
draw_mask(draw, points, shape_type, label, 255 - label)
|
||
|
||
mask = np.array(mask)
|
||
mask = Image.fromarray(mask)
|
||
mask.save(str(Path(mask_saved_path) / (str(image_path).split(".")[0] + ".png")))
|
||
|
||
os.remove(json_path)
|
||
|
||
"""
|
||
将root_path下labelme生成的json文件全部进行处理：
|
||
1. 有原图匹配的json文件，会转换成mask，存储到mask_saved_path路径下
|
||
2. 没有原图，但是有json文件的，直接删除该json文件
|
||
3. 有原图，但是没有json文件的，会在mask_saved_path下生成一个纯黑背景图片
|
||
root_path: 存储着原图和json文件，原图后缀名尽量为.jpg
|
||
"""
|
||
def convert_labelme_jsons2masks(classes, root_path: str, mask_saved_path: str, original_image_suffix=".jpg"):
|
||
assert 0 < len(classes) <= 128
|
||
original_images = set(
|
||
map(
|
||
lambda name: str(name).split(".")[0],
|
||
Path(root_path).glob(pattern=f"*{original_image_suffix}")
|
||
)
|
||
)
|
||
json_files = Path(root_path).glob(pattern="*.json")
|
||
for json_file in json_files:
|
||
filename = str(json_file).split(".")[0]
|
||
if filename in original_images:
|
||
labelme_json2mask(classes, str(json_file), mask_saved_path)
|
||
original_images.remove(filename)
|
||
else:
|
||
os.remove(json_file)
|
||
|
||
if len(original_images) != 0:
|
||
for image_filename in original_images:
|
||
image_path = image_filename + f"{original_image_suffix}"
|
||
image = Image.open(image_path)
|
||
height, width = image.height, image.width
|
||
image.close()
|
||
mask = np.zeros((height, width), dtype=np.uint8)
|
||
mask = Image.fromarray(mask)
|
||
mask.save(str(Path(mask_saved_path) / (os.path.basename(image_filename) + ".png")))
|
||
|
||
"""
|
||
将混淆矩阵得到的尺度(scores)组合成表格形式输出到控制台
|
||
scores: 混淆矩阵的尺度(scores)
|
||
"""
|
||
def confusion_matrix_scores2table(scores):
|
||
assert scores is not None and isinstance(scores, dict)
|
||
|
||
classes = [tp[0] for tp in scores["classes_precision"]]
|
||
cls_precision_list = [tp[-1] for tp in scores["classes_precision"]]
|
||
cls_recall_list = [tp[-1] for tp in scores["classes_recall"]]
|
||
cls_iou_list = [tp[-1] for tp in scores["classes_iou"]]
|
||
table1 = tabulate(
|
||
tabular_data=np.concatenate(
|
||
(
|
||
np.asarray(classes).reshape(-1, 1),
|
||
np.asarray(cls_precision_list).reshape(-1, 1),
|
||
np.asarray(cls_recall_list).reshape(-1, 1),
|
||
np.asarray(cls_iou_list).reshape(-1, 1)
|
||
), 1
|
||
),
|
||
headers=["classes", "precision", "recall", "iou"],
|
||
tablefmt="grid"
|
||
)
|
||
|
||
avg_precision = scores["avg_precision"]
|
||
avg_recall = scores["avg_recall"]
|
||
avg_iou = scores["avg_iou"]
|
||
accuracy = scores["accuracy"]
|
||
table2 = tabulate(
|
||
tabular_data=[(avg_precision, avg_recall, avg_iou, accuracy)],
|
||
headers=["avg_precision", "avg_recall", "avg_iou", "accuracy"],
|
||
tablefmt="grid"
|
||
)
|
||
|
||
table = tabulate(
|
||
tabular_data=np.concatenate(
|
||
(
|
||
np.asarray(["single", "overall"]).reshape(-1, 1),
|
||
np.asarray([table1, table2]).reshape(-1, 1)
|
||
), 1
|
||
),
|
||
headers=["table type", "table"],
|
||
tablefmt="grid"
|
||
)
|
||
|
||
return table
|
||
|
||
|
||
"""
|
||
相加混淆矩阵得到的两个scores
|
||
|
||
返回值：
|
||
相加后的混淆矩阵
|
||
"""
|
||
def sum_2_confusion_matrix_scores(scores_left: dict, scores_right: dict):
|
||
scores_left["classes_precision"] = [
|
||
(tp[0][0], tp[0][-1] + tp[-1][-1]) for tp in zip(scores_left["classes_precision"], scores_right["classes_precision"])
|
||
]
|
||
scores_left["classes_recall"] = [
|
||
(tp[0][0], tp[0][-1] + tp[-1][-1]) for tp in zip(scores_left["classes_recall"], scores_right["classes_recall"])
|
||
]
|
||
scores_left["classes_iou"] = [
|
||
(tp[0][0], tp[0][-1] + tp[-1][-1]) for tp in zip(scores_left["classes_iou"], scores_right["classes_iou"])
|
||
]
|
||
|
||
scores_left["avg_precision"] = scores_left["avg_precision"] + scores_right["avg_precision"]
|
||
scores_left["avg_recall"] = scores_left["avg_recall"] + scores_right["avg_recall"]
|
||
scores_left["avg_iou"] = scores_left["avg_iou"] + scores_right["avg_iou"]
|
||
scores_left["accuracy"] = scores_left["accuracy"] + scores_right["accuracy"]
|
||
|
||
return scores_left
|
||
|
||
"""
|
||
将混淆矩阵列表内的scores进行相加
|
||
@:param scores_list: 得分列表
|
||
@:return 相加后的得分
|
||
"""
|
||
def sum_confusion_matrix_scores_list(scores_list):
|
||
if len(scores_list) == 1:
|
||
return scores_list[0]
|
||
|
||
result = scores_list[0]
|
||
for i in range(1, len(scores_list)):
|
||
result = sum_2_confusion_matrix_scores(result, scores_list[i])
|
||
return result
|
||
|
||
"""
|
||
对混淆矩阵得出的scores_list相加后求平均
|
||
|
||
返回值：
|
||
相加后求平均的scores
|
||
"""
|
||
def avg_confusion_matrix_scores_list(scores_list):
|
||
assert scores_list is not None and len(scores_list) >= 1
|
||
result = sum_confusion_matrix_scores_list(scores_list)
|
||
|
||
result["classes_precision"] = [
|
||
(tp[0], tp[-1] / len(scores_list)) for tp in result["classes_precision"]
|
||
]
|
||
result["classes_recall"] = [
|
||
(tp[0], tp[-1] / len(scores_list)) for tp in result["classes_recall"]
|
||
]
|
||
result["classes_iou"] = [
|
||
(tp[0], tp[-1] / len(scores_list)) for tp in result["classes_iou"]
|
||
]
|
||
|
||
result["avg_precision"] = result["avg_precision"] / len(scores_list)
|
||
result["avg_recall"] = result["avg_recall"] / len(scores_list)
|
||
result["avg_iou"] = result["avg_iou"] / len(scores_list)
|
||
result["accuracy"] = result["accuracy"] / len(scores_list)
|
||
|
||
return result
|
||
|
||
"""
|
||
对原始作为x的输入图像进行增强预处理，产生相同大小的图片(旋转、翻转、亮度调整)
|
||
ts是pytorch工具包，经过该工具包处理后图像如果和原本的不同，
|
||
就会保存在磁盘上，以达到增强数据的目的，请先执行该函数之后，再对原始数
|
||
据图像进行人工标注。
|
||
root_path目录下的数据只有图片，且图片后缀名一致
|
||
|
||
root_path: 作为x的原始输入图像所在目录
|
||
ts: 预处理策略
|
||
"""
|
||
def augment_raw_images2(
|
||
root_path,
|
||
ts=transforms.Compose(
|
||
[
|
||
transforms.RandomHorizontalFlip(),
|
||
transforms.RandomVerticalFlip(),
|
||
transforms.RandomRotation(degrees=30),
|
||
transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3)
|
||
]
|
||
)
|
||
):
|
||
image_paths = Path(root_path).glob(pattern="*")
|
||
for image_path in image_paths:
|
||
counter = 0
|
||
image_filename, image_suffix = os.path.splitext(image_path)
|
||
|
||
image = Image.open(image_path)
|
||
image_np = np.array(image)
|
||
|
||
|
||
for transform in ts.transforms:
|
||
new_image = transform(Image.fromarray(image_np))
|
||
new_image_np = np.array(new_image)
|
||
|
||
if not np.array_equal(image_np, new_image_np):
|
||
new_image_copy = Image.fromarray(new_image_np)
|
||
new_image_copy.save(str(Path(f"{image_filename}_{counter}{image_suffix}")))
|
||
new_image_copy.close()
|
||
counter += 1
|
||
|
||
new_image.close()
|
||
|
||
image.close()
|
||
|
||
|
||
|
||
"""
|
||
对原始作为x的输入图像进行增强预处理，产生image_cropped_shape大小的图片
|
||
现将图像resize为image_resized_shape大小，然后进行1次裁剪和1次随机裁剪，裁剪的图像保留下来，原始图像放入to_path中
|
||
ts是pytorch工具包，经过该工具包处理后图像如果和原本的不同，
|
||
就会保存在磁盘上，以达到增强数据的目的，请先执行该函数之后，再对原始数
|
||
据图像进行人工标注。
|
||
from_path目录下的数据只有图片，且图片后缀名一致
|
||
|
||
from_path: 作为x的原始输入图像所在目录
|
||
to_path: 处理后的原始图像放入哪里，如果为None，就删除原始图像
|
||
image_resized_shape: 图像resize之后的大小, image_cropped_shape每个维度必须小于image_resized_shape
|
||
image_cropped_shape: 图像裁剪后的大小，image_cropped_shape每个维度必须小于image_resized_shape
|
||
ts: 预处理策略
|
||
"""
|
||
def augment_raw_images(
|
||
from_path,
|
||
to_path="to/path",
|
||
image_resized_shape=(256, 256),
|
||
image_cropped_shape=(224, 224),
|
||
ts=None
|
||
):
|
||
if ts is None:
|
||
ts = transforms.Compose(
|
||
[
|
||
transforms.Resize(image_resized_shape, interpolation=InterpolationMode.BILINEAR),
|
||
transforms.RandomCrop(image_cropped_shape),
|
||
transforms.RandomResizedCrop(image_cropped_shape)
|
||
]
|
||
)
|
||
image_paths = Path(from_path).glob("*")
|
||
for image_path in image_paths:
|
||
counter = 0
|
||
image_filename, image_suffix = os.path.splitext(image_path)
|
||
with Image.open(image_path) as image:
|
||
image = ts.transforms[0](image)
|
||
image_copy_np = copy.deepcopy(np.array(image))
|
||
for transform in ts.transforms[0:]:
|
||
image = transform(image)
|
||
image_np = np.array(image)
|
||
if not np.array_equal(image_np, image_copy_np):
|
||
image.save(str(Path(f"{image_filename}_{counter}{image_suffix}")))
|
||
counter = counter + 1
|
||
image.close()
|
||
image = Image.fromarray(image_copy_np)
|
||
if to_path:
|
||
Path(image_path).rename(Path(to_path) / f"{os.path.basename(image_path)}")
|
||
else:
|
||
Path(image_path).unlink()
|
||
|
||
|
||
"""
|
||
对验证数据集中的图片进行大小的统一，以便其拥有统一的大小，可以进行批次训练
|
||
from_path: 验证数据集所在的目录
|
||
to_path: 原始数据应该转移到哪里
|
||
resized_shape: (height, width), resize后的大小
|
||
"""
|
||
def resize_val_images(from_path, to_path, resized_shape):
|
||
image_paths = Path(from_path).glob(pattern="*")
|
||
for image_path in image_paths:
|
||
original_image = Image.open(image_path)
|
||
original_image_np = np.array(original_image)
|
||
resized_image = Image.fromarray(original_image_np).resize(resized_shape)
|
||
original_image.close()
|
||
|
||
if not to_path:
|
||
Path(image_path).unlink(missing_ok=True)
|
||
else:
|
||
Path(image_path).rename(Path(to_path) / os.path.basename(image_path))
|
||
|
||
resized_image.save(image_path)
|
||
resized_image.close()
|
||
|
||
|
||
"""
|
||
将一张图片按照尺寸裁剪为多张图片
|
||
@:param image: 图片
|
||
@:param crop_size: 裁剪尺寸，为tuple(image_height, image_width)
|
||
|
||
@:return 裁剪之后的图片列表
|
||
"""
|
||
def crop_image2images(image: Image, crop_size):
|
||
image_np = np.array(image)
|
||
image_height, image_width = image_np.shape[:-1]
|
||
left_image_height, left_image_width = image_np.shape[:-1]
|
||
crop_height, crop_width = crop_size
|
||
left_upper = (0, 0)
|
||
right_lower = (crop_width, crop_height)
|
||
image_list = []
|
||
|
||
while left_image_width / crop_width >= 1 or left_image_height / crop_height >= 1:
|
||
if left_image_width / crop_width >= 1 and left_image_height / crop_height >= 1:
|
||
new_image = image.crop((*left_upper, *right_lower))
|
||
left_image_width -= crop_width
|
||
left_upper = (left_upper[0] + crop_width, left_upper[-1])
|
||
right_lower = (right_lower[0] + crop_width, right_lower[-1])
|
||
image_list.append(new_image)
|
||
elif left_image_height / crop_height >= 1:
|
||
left_image_width = image_width
|
||
left_image_height -= crop_height
|
||
left_upper = (0, image_height - left_image_height)
|
||
right_lower = (crop_width, image_height - left_image_height + crop_height)
|
||
else:
|
||
break
|
||
return image_list
|
||
|
||
"""
|
||
将目录下的所有图片进行裁剪
|
||
@:param root_path: 图片的目录
|
||
@:param to: 原图片应该转移到哪里
|
||
@:param crop_size: 裁剪大小, tuple(crop_height, crop_width)
|
||
"""
|
||
def crop_images2small_images(root_path, to, crop_size):
|
||
image_paths = Path(root_path).glob(pattern="*")
|
||
for image_path in image_paths:
|
||
image = Image.open(image_path)
|
||
image_cropped_list = crop_image2images(image, crop_size)
|
||
for idx, image_cropped in enumerate(image_cropped_list):
|
||
image_cropped.save(
|
||
f"_{idx}".join(os.path.splitext(image_path))
|
||
)
|
||
image_cropped.close()
|
||
image.close()
|
||
if to is None:
|
||
Path(image_path).unlink(missing_ok=True)
|
||
else:
|
||
Path(image_path).rename(
|
||
str(
|
||
Path(to) / os.path.basename(image_path)
|
||
)
|
||
)
|
||
|
||
"""
|
||
判断是否能够多gpu分布式并行运算
|
||
"""
|
||
def distributed_enabled():
|
||
return torch.cuda.is_available() and torch.cuda.device_count() > 1 and torch.__version__ >= "0.4.0"
|
||
|
||
if __name__ == "__main__":
|
||
# crop_images2small_images(
|
||
# root_path="dataset/train/images",
|
||
# to=None,
|
||
# crop_size=(512, 512)
|
||
# )
|
||
|
||
|
||
# augment_raw_images2(root_path="dataset/train/images")
|
||
|
||
crop_images2small_images(
|
||
root_path="dataset/test",
|
||
to=None,
|
||
crop_size=(512, 512)
|
||
)
|
||
|
||
# augment_raw_images2(root_path="dataset/val/images")
|
||
|
||
# resize_val_images(
|
||
# from_path="dataset/test",
|
||
# to_path=None,
|
||
# resized_shape=(1024, 1024)
|
||
# )
|
||
|
||
# convert_labelme_jsons2masks(
|
||
# classes=[
|
||
# "background",
|
||
# "leaf"
|
||
# ],
|
||
# root_path="dataset/train/images",
|
||
# mask_saved_path="dataset/train/labels"
|
||
# ) |