干眼—图像增强算法 扭曲+噪声

# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""
import PIL.Image as IM
import cv2
import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage.interpolation import map_coordinates
from scipy.ndimage.filters import gaussian_filter
import time
RADIUS_LIST = [50,84,126,166,208,247,287,328,378,428,486]
HU_LENGTH = [16,24,24,24,24,24,24,24,29,36,40]
CIRCLE_NUMS = 11
#随机点破裂仿射
# Function to distort image
def elastic_transform(image, alpha, sigma, alpha_affine, random_state=None):
    """Elastic deformation of images as described in [Simard2003]_ (with modifications).
    .. [Simard2003] Simard, Steinkraus and Platt, "Best Practices for
         Convolutional Neural Networks applied to Visual Document Analysis", in
         Proc. of the International Conference on Document Analysis and
         Recognition, 2003.

     Based on https://gist.github.com/erniejunior/601cdf56d2b424757de5
    """
    if random_state is None:
        random_state = np.random.RandomState(None)

    shape = image.shape
    shape_size = shape[:2]
    
    # Random affine
    center_square = np.float32(shape_size) // 2
    square_size = min(shape_size) // 3
    pts1 = np.float32([center_square + square_size, [center_square[0]+square_size, center_square[1]-square_size], center_square - square_size])
    pts2 = pts1 + random_state.uniform(-alpha_affine, alpha_affine, size=pts1.shape).astype(np.float32)
    M = cv2.getAffineTransform(pts1, pts2)
    image = cv2.warpAffine(image, M, shape_size[::-1], borderMode=cv2.BORDER_REFLECT_101)

    dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma) * alpha
    dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma) * alpha
    dz = np.zeros_like(dx)

    x, y, z = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]), np.arange(shape[2]))
    indices = np.reshape(y+dy, (-1, 1)), np.reshape(x+dx, (-1, 1)), np.reshape(z, (-1, 1))

    return map_coordinates(image, indices, order=1, mode='reflect').reshape(shape)

#加噪声
def noise(img,nums=100):
    
    img_black_wai = np.zeros((1000,1000)).astype(np.uint8)
    img_black_wai = cv2.circle(img_black_wai,(500,500),RADIUS_LIST[-1],(255,0,0),-1)
    pos_black_list = list(zip(np.where(img_black_wai==0)[0],np.where(img_black_wai==0)[1]))
    random_pos_list = np.random.randint(0,len(pos_black_list),np.random.randint(30,120))
    for pos in random_pos_list:
        pos = pos_black_list[pos]
        w = np.random.randint(0,50)
        h = np.random.randint(0,50)
#        y = np.random.randint(50,img.shape[0]-50)
#        x = np.random.randint(50,img.shape[1]-50)
        xmin = max(0,pos[1]-w)
        ymin = max(0,pos[0]-h)
        xmax = min(img.shape[1],pos[1]+w)
        ymax = min(img.shape[0],pos[0]+h)
        img[ymin:ymax,xmin:xmax] = 255
    for i in range(np.random.randint(3000,15000)):
        y = np.random.randint(0,img.shape[0])
        x = np.random.randint(0,img.shape[1])
        img[y,x]=255
    return img
# Define function to draw a grid
def draw_grid(im, grid_size):
    # Draw grid lines
    for i in range(0, im.shape[1], grid_size):
        cv2.line(im, (i, 0), (i, im.shape[0]), color=(255,))
    for j in range(0, im.shape[0], grid_size):
        cv2.line(im, (0, j), (im.shape[1], j), color=(255,))
        
def get_image():        
    
    img_name = 'Affine_Line_'+str(time.time()).replace('.','')+'_xkk'+str(np.random.randint(0,92395))+'.jpg'
    img_black = np.zeros((1000,1000)).astype(np.uint8)
    plt.imshow(img_black,cmap='gray')
    
    #img_black = cv2.circle(img_black,(200,200),3,(255,0,0),1)
    for i in range(0,CIRCLE_NUMS):    
        img_black = cv2.circle(img_black,(500,500),RADIUS_LIST[i]+np.random.randint(-2,2),(255,0,0),HU_LENGTH[i]+np.random.randint(-2,2))
    #cv2.imshow('img_black',img_black)
    plt.imshow(img_black,cmap='gray')
    plt.show()
    #IM.fromarray(img_black).show()
    
    img_canny = cv2.Canny(img_black,50,200)
    #img_canny 
    plt.imshow(img_canny,cmap='gray')
    plt.show()
    
    pos_edge = list(zip(np.where(img_canny==255)[0],np.where(img_canny==255)[1]))
    
    random_pos = np.random.randint(0,len(pos_edge),10)
    
    #pos = pos_edge[random_pos[0]]
    #生成干扰圆
    for idx in random_pos:
        pos = pos_edge[idx]
        img_black = cv2.circle(img_black,(pos[1],pos[0]),20,(255,0,0),-1)
    plt.imshow(img_black,cmap='gray')
    plt.show()
    
    #for x in range(0,1000):
    #找最外层的圆 边界点，来画睫毛
    
    img_black_wai = np.zeros((1000,1000)).astype(np.uint8)
    img_black_wai = cv2.circle(img_black_wai,(500,500),RADIUS_LIST[-1],(255,0,0),1)
    img_black_wai = cv2.Canny(img_black_wai,50,200)
    pos_edge_wai_list = list(zip(np.where(img_canny==255)[0],np.where(img_canny==255)[1]))
    
    count=0
    #上方睫毛
    for line_start in pos_edge_wai_list:
        count+=1
        if line_start[0]>30 or count%100==0:
            continue
    #    print(1)
        lengthx = np.random.randint(-200,300)
        lengthy = np.random.randint(0,200)
        img_black = cv2.line(img_black,(line_start[1],line_start[0]),(line_start[1]-lengthx ,line_start[0]+lengthy ),(0,0,0),5)
        
    for line_start in pos_edge_wai_list:
        count+=1
        if line_start[0]>220 or count%70!=0:
            continue
    #    print(1)
        lengthx = np.random.randint(-300,400)
        lengthy = np.random.randint(0,200)
        img_black = cv2.line(img_black,(line_start[1],line_start[0]),(line_start[1]-lengthx ,line_start[0]+lengthy ),(0,0,0),5)
    
    #下方睫毛
    for line_start in pos_edge_wai_list:
        count+=1
        if line_start[0]<800 or count%200!=0:
            continue
    #    print(1)
        lengthx = np.random.randint(-50,50)
        lengthy = np.random.randint(0,100)
        img_black = cv2.line(img_black,(line_start[1],line_start[0]),(line_start[1]-lengthx ,30+line_start[0]-lengthy ),(0,0,0),5)
    plt.imshow(img_black,cmap='gray')
    plt.show()
    
    
    
    img_black = cv2.cvtColor(img_black,cv2.COLOR_GRAY2RGB)
    count = 0
    random_idx = np.random.randint(0,len(pos_edge),np.random.randint(2,18))
    affine_pos = []
    for idx in random_idx :
        pos = pos_edge[idx]
        count+=1 
    #    if count%5000!=0:
    #        continue
        if pos[0]<50 or pos[0]>img_black.shape[0]-50:
            continue
        if pos[1]<50 or pos[1]>img_black.shape[1]-50:
            continue
        if pos[0]<300:
            continue
        im_merge = img_black[pos[0]-50:pos[0]+50,pos[1]-50:pos[1]+50,:]
        im_merge_t = elastic_transform(im_merge, im_merge.shape[1] * np.random.randint(5,12), im_merge.shape[1] * 0.1*np.random.randint(1,5), im_merge.shape[1] * 0.02*np.random.randint(2,5))
        img_black[pos[0]-50:pos[0]+50,pos[1]-50:pos[1]+50] = im_merge_t
        savepath = 'D:\\xkkAI\\DryEye\\keras-frcnn-master\\xkk0715\\train_data\\'+img_name
        newpos = pos+(savepath,)
        affine_pos.append(newpos)
    savepath = 'D:\\xkkAI\\DryEye\\keras-frcnn-master\\xkk0715\\train_data\\'+img_name
    img_black = noise(img_black)
    IM.fromarray(img_black).save(savepath )
    #print(str(time.time()))
    for pos in affine_pos:
        xmin = pos[1]-50
        print(pos)
        ymin = pos[0]-50
        xmax = pos[1]+50
        ymax = pos[0]+50
        img_final = cv2.rectangle(img_black,(xmin,ymin),(xmax,ymax),(255,0,0),3)
        
    plt.imshow(img_final)
    plt.show()
#    plt.clf()
    return affine_pos

with open(r'D:\xkkAI\DryEye\keras-frcnn-master\xkk0715\anno.txt','a') as f:
    for i in range(5000):    
        affine_pos = get_image()
        for pos in affine_pos:        
            xmin = pos[1]-50
            print(pos)
            ymin = pos[0]-50
            xmax = pos[1]+50
            ymax = pos[0]+50
#            img_final = cv2.rectangle(img_black,(xmin,ymin),(xmax,ymax),(255,0,0),3)
            f.write(pos[2]+','+str(xmin)+','+str(ymin)+','+str(xmax)+','+str(ymax)+',posun\n')
#IM.fromarray(img_final).show()