def circle2rectangle(path):
img = Image.open(path)
center, radius = get_radius_center(img)
img_np = np.array(img)
rec = []
r1 = 220
r2 = radius
for i in np.linspace(0, 2 * math.pi, 448):
col = []
for j in np.linspace(r1, r2, 448):
x = int(j * math.cos(i) + center[0])
y = int(j * math.sin(i) + center[1])
col.append(img_np[y][x])
rec.append(col)
rec = np.array(rec)
target_rec = cv2.cvtColor(rec, cv2.COLOR_GRAY2RGB)
target_rec_pillow = Image.fromarray(target_rec)
target_rec_pillow = resize(target_rec_pillow, 448)
# img_np = np.array(img)
# img_np = np.moveaxis(img_np, -1, 0)
# img_np = img_np[np.newaxis, ...]
# img_tensor = torch.tensor(img_np, dtype=torch.float32)
img_tensor = to_tensor(target_rec_pillow)
img_tensor = img_tensor.numpy()
img_tensor = img_tensor[np.newaxis, ...]
img_tensor = torch.from_numpy(img_tensor)
return img_tensor, target_rec_pillow
def img2tensor(image_path):
# 读取图片
image_ori = Image.open(image_path)
# 获取目标位置
center, radius = get_radius_center(image_ori)
left_up_corner = np.array(center) - radius
image = crop(image_ori, left_up_corner[1], left_up_corner[0], radius * 2,
radius * 2)
# 添加黑色背景
image = add_black_background(image)
# 缩放
image = image.resize((512, 512))
# to tensor
image = np.array(image)
if len(image.shape) == 2:
image = image[:, :, np.newaxis]
# swap color axis because
# numpy image: H x W x C
# torch image: C X H X W
image = image.transpose((2, 0, 1))
if image.shape[0] == 1:
image = np.repeat(image, 3, axis=0)
image = image[np.newaxis, ...]
image = torch.from_numpy(image).float()
return image, image_ori
def img2tensor(path):
img = Image.open(path)
center, radius = get_radius_center(img)
center_x = center[0]
center_y = center[1]
center_r = radius
half_width = int(center_r * 1.01) # height = width
img = img.crop((int(center_x - half_width), int(center_y - half_width),
int(center_x + half_width), int(center_y + half_width)))
# 生成mask覆盖非检测区域
size = img.size
circle = np.zeros(size, dtype='uint8') # 黑色背景
cv2.circle(circle, (size[0] // 2, size[1] // 2), size[0] // 2, 1,
-1) # 中心圆内不改变
circle = np.stack((circle, ) * 3, -1) # 扩展维度 gray to rgb
img_np = cv2.cvtColor(np.array(img),
cv2.COLOR_GRAY2BGR) # 黑白照片其实不需要RGB2BGR
img_np = img_np * circle # 0去除1留存
img = Image.fromarray(np.array(cv2.cvtColor(
img_np, cv2.COLOR_BGR2RGB))) # 黑白照片其实不需要RGB2BGR
img = resize(img, 448)
# img_np = np.array(img)
# img_np = np.moveaxis(img_np, -1, 0)
# img_np = img_np[np.newaxis, ...]
# img_tensor = torch.tensor(img_np, dtype=torch.float32)
img_tensor = to_tensor(img)
img_tensor = img_tensor.numpy()
img_tensor = img_tensor[np.newaxis, ...]
img_tensor = torch.from_numpy(img_tensor)
return img_tensor, img
def __call__(self, sample):
image, label = sample['image'], sample['label']
# print(sample['json_path'])
if random.random() < self.p:
if sample['json_path'][-4:] != 'None':
image = generate_new_data_v2(sample['image_path'],
sample['json_path'])
else:
center, radius = get_radius_center(image)
left_up_corner = np.array(center) - radius
image = crop(image, left_up_corner[1], left_up_corner[0],
radius * 2, radius * 2)
else:
center, radius = get_radius_center(image)
left_up_corner = np.array(center) - radius
image = crop(image, left_up_corner[1], left_up_corner[0],
radius * 2, radius * 2)
image = image.resize((1024, 1024))
return {'image': image, 'label': label}
def get_random_good_as_background():
root_path = os.path.join('data_all', 'side', 'train', 'good')
images = os.listdir(root_path)
random_image = random.choice(images)
random_image = cv2.imread(os.path.join(root_path, random_image))
# 获取工件中心坐标和半径
center, radius = get_radius_center(Image.fromarray(random_image))
# 获取原图裁切目标图 带有损伤信息 cv2 crop [y:y+h,x:x+h]
random_image = random_image[center[1] - radius:center[1] + radius,
center[0] - radius:center[0] + radius]
return Image.fromarray(random_image)
def get_background(image, label):
# 返回裁切后的正方形工件背景
# 获取工件中心坐标和半径
center, radius = get_radius_center(Image.fromarray(image))
left_up_corner = np.array(center) - radius
# 获取原图裁切目标图 带有损伤信息 cv2 crop [y:y+h,x:x+h]
origin_crop_image = image[center[1] - radius:center[1] + radius,
center[0] - radius:center[0] + radius]
# 根据labels创建mask遮罩
mask = np.zeros((origin_crop_image.shape[0], origin_crop_image.shape[1]),
np.uint8)
# 根据圆心坐标移动labels坐标 labels [np.array,...]
for i in range(len(label)):
label[i] = label[i] - np.repeat(
[left_up_corner], len(label[i]), axis=0)
cv2.fillPoly(mask, [label[i]], 255)
# inpaint
# Image.fromarray(mask).show()
background = cv2.inpaint(origin_crop_image, mask, 3, cv2.INPAINT_NS)
return Image.fromarray(background)
def generate_new_data(image_path, labels_path):
image_cv2 = cv2.imread(
image_path) # cv2 load image BGR image has 3 channels
with open(labels_path) as f: # json load labels
json_labels = json.load(f)
# 解析labels
labels = [] # 损伤位置组成的list [np.array,...]
for i in json_labels['shapes']:
labels.append(np.array(i['points'], np.int32))
# 获取工件中心坐标和半径
center, radius = get_radius_center(Image.fromarray(image_cv2))
left_up_corner = np.array(center) - radius
# 获取原图裁切目标图 带有损伤信息 cv2 crop [y:y+h,x:x+h]
origin_crop_image = image_cv2[center[1] - radius:center[1] + radius,
center[0] - radius:center[0] + radius]
# 根据labels创建mask遮罩
mask = np.zeros((origin_crop_image.shape[0], origin_crop_image.shape[1]),
np.uint8)
# 根据圆心坐标移动labels坐标 labels [np.array,...]
for i in range(len(labels)):
labels[i] = labels[i] - np.repeat(
[left_up_corner], len(labels[i]), axis=0)
cv2.fillPoly(mask, [labels[i]], 255)
# inpaint
background = cv2.inpaint(origin_crop_image, mask, 3, cv2.INPAINT_NS)
# Image.fromarray(background).show()
# 分别调整每个损伤
new_image = copy.deepcopy(background)
new_image = cv2.cvtColor(new_image, cv2.COLOR_BGR2GRAY)
origin_copy = copy.deepcopy(origin_crop_image)
origin_copy = cv2.cvtColor(origin_copy, cv2.COLOR_BGR2GRAY)
for i in labels:
# 只要损伤区域不要其他区域
single_mask = np.zeros(new_image.shape, np.uint8)
single_mask = cv2.fillPoly(single_mask, [i], 1)
single_mask = np.multiply(single_mask, origin_copy)
# Image.fromarray(single_mask).show()
# 获取随机移动值
xc = radius # 中心坐标
yc = radius # 中心坐标
x0 = (np.min(i[:, 0]) + np.max(i[:, 0])) / 2 # 损伤位置 平均中心坐标
y0 = (np.min(i[:, 1]) + np.max(i[:, 1])) / 2 # 损伤位置 平均中心坐标
k = (yc - y0) / (xc - x0) # 移动线路 损伤与工件圆心的连线
b = yc - k * xc
delta_x = int((random.random() - 0.5) * 50)
if math.isnan(k * (x0 + delta_x) + b - y0):
continue
delta_y = int(k * (x0 + delta_x) + b - y0)
if 1:
# 移动损伤位置
(mx, my) = np.meshgrid(np.arange(single_mask.shape[1]),
np.arange(single_mask.shape[0]))
ox = (mx - delta_x).astype(np.float32)
oy = (my - delta_y).astype(np.float32)
single_mask = cv2.remap(single_mask, ox, oy, cv2.INTER_LINEAR)
else:
delta_x, delta_y = 0, 0
# 对应的改变背景留黑位置
adjust_i = i + np.array([delta_x, delta_y])
# 只要其他区域不要损伤区域
background_mask = np.ones(new_image.shape, np.uint8)
background_mask = cv2.fillPoly(background_mask, [adjust_i], 0)
background_mask = np.multiply(background_mask, new_image)
new_image = np.add(single_mask, background_mask)
# 旋转图像
new_image = Image.fromarray(new_image)
new_image = new_image.rotate(random.random() * 360,
resample=Image.BICUBIC)
new_image = np.array(new_image)
return Image.fromarray(new_image)