Top.Mail.Ru

Python. Машинное обучение. Продолжение 2.

Продолжим изучать машинное обучение на Python. На прошлом уроке мы написали модель поиска дефектов. Теперь используя эту модель, найдём дефект и заменим его на нормальный пиксель при помощи модели, разработанной на уроке Python. Машинное обучение. Напомню кратко модель: мы в обучающую выборку загоняем окружения дефектных точек, точнее, цвета тех точек, которые находятся вокруг данной точки. Напомню, что цвет точки – это числа от 0 до 255, для каждой составляющих RGB (Red, Blue, Green). Например, если размер скользящего окна  равен 5, то входной элемент состоит из 24 точек (сама эта точке в него не входит). Соответственно, на вход подается 72 числа. На выходе – одно число, это составляющая R, G или B анализируемой точки. Соответственно, нам потребуется три модели машинного обучения, для каждого из цветов RGB, которые мы будет обучать и использовать отдельно.

Итак, вот код программы на Python:

%%time

import random
import pandas as pd
import numpy as np
import datetime
from sklearn.ensemble import RandomForestClassifier
from PIL import Image, ImageDraw #Подключим библиотеки для работы с картинками

def CreateModel(selection, out_name):
    ls=["RY", "GY", "BY"]
    ls.remove(out_name)
    Y = selection[out_name]
    X = selection.drop(ls, axis=1)
    model = RandomForestClassifier()
    model.fit(X, Y)
    return model

def GetColor(selection, ls, model):
    X = selection.drop(ls, axis=1)
    predicted = model.predict(X)
    return predicted

def ComparePoints(point1, point2):
    if point1[0]!=point2[0]:
        return False
    if point1[1]!=point2[1]:
        return False
    if point1[2]!=point2[2]:
        return False    
    return True

#Создаем словарь строки выборки
def CreateDicDefected(pix, i, j, W, result=-1):
    dic={}
    k=1            
    for x in range(i-W, i+W+1):
        for y in range(j-W,j+W+1):
            if x==i and y==j:
                continue
            dic["R"+str(k)] = pix[x, y][0]
            dic["G"+str(k)] = pix[x, y][1]
            dic["B"+str(k)] = pix[x, y][2]
            k=k+1
    if result==-1:
        dic["RY"]=pix[i, j][0]
        dic["GY"]=pix[i, j][1]
        dic["BY"]=pix[i, j][2]
    else:
        dic["Y"]=result
    return dic
            
#Создание выборки для поиска дефектов
def CreateDefectionSelection(FileOriginal, FileDefected, W):
    imageOriginal = Image.open(FileOriginal)
    widthOriginal = imageOriginal.size[0] #Определяем ширину. 
    heightOriginal = imageOriginal.size[1] #Определяем высоту.         
    imageDefected = Image.open(FileDefected)
    widthDefected = imageDefected.size[0] #Определяем ширину. 
    heightDefected = imageDefected.size[1] #Определяем высоту.
    pixOriginal = imageOriginal.load() #Выгружаем значения пикселей.
    pixDefected = imageDefected.load() #Выгружаем значения пикселей.
    count=widthOriginal*heightOriginal
    if widthOriginal!=widthDefected:
        print("Не равна ширина оригинала и испорченной картинки")
        return NaN
    if heightOriginal!=heightDefected:
        print("Не равна высота оригинала и испорченной картинки")
        return NaN  
    ls=[]
    for i in range(W,widthOriginal-W):
        for j in range(W,heightOriginal-W):      
            if ComparePoints(pixOriginal[i,j],pixDefected[i,j]):
                if random.randint(1,int(count/2000))==1:
                    ls.append(CreateDicDefected(pixOriginal, i, j, W, 1))
            else:
                ls.append(CreateDicDefected(pixDefected, i, j, W, 0))
    return pd.DataFrame(ls)
            
def CreateRepairSelection(FileName, W):
    image = Image.open(FileName) #Открываем изображение. 
    draw = ImageDraw.Draw(image) #Создаем инструмент для рисования. 
    width = image.size[0] #Определяем ширину. 
    height = image.size[1] #Определяем высоту. 	
    pix = image.load() #Выгружаем значения пикселей.
    count=width*height
    ls_train=[]
    for i in range(W,width-W):
        for j in range(W,height-W):        
            if random.randint(1,int(count/2000))==1:
                dic={}
                k=1            
                for x in range(i-W, i+W+1):
                    for y in range(j-W,j+W+1):
                        if x==i and y==j:
                            continue
                        dic["R"+str(k)] = pix[x, y][0]
                        dic["G"+str(k)] = pix[x, y][1]
                        dic["B"+str(k)] = pix[x, y][2]
                        k=k+1
                dic["RY"]=pix[i, j][0]
                dic["GY"]=pix[i, j][1]
                dic["BY"]=pix[i, j][2]
                ls_train.append(dic)
    return ls_train

W=2 # Полуразмер окна

#Создаем модель дефектов
defect=CreateDefectionSelection("d:\\rep\\origtrain.png","d:\\rep\\origtrainspoil.png",W)
data_train_defect=pd.DataFrame(defect)
Y_defect = data_train_defect["Y"]
X_defect = data_train_defect.drop(["Y"], axis=1)
model_defect = RandomForestClassifier()
model_defect.fit(X_defect,Y_defect)

#Создаем модель восстановления
rep=CreateRepairSelection("d:\\rep\\origtrainspoil.png",W)
data_train_rep=pd.DataFrame(rep)
modelR=CreateModel(data_train_rep, "RY")
modelG=CreateModel(data_train_rep, "GY")
modelB=CreateModel(data_train_rep, "BY") 

image = Image.open("d:\\rep\\origtestspoild.png") #Открываем изображение. 
draw = ImageDraw.Draw(image) #Создаем инструмент для рисования. 
width = image.size[0] #Определяем ширину. 
height = image.size[1] #Определяем высоту. 	
pix = image.load() #Выгружаем значения пикселей.
count=width*height

new_image = Image.new("RGB", (width,height), (0,0,0))
draw = ImageDraw.Draw(new_image)  # Создаем инструмент для рисования
count=1
limit=999999999
ls_defects=[]
for i in range(W,width-W):
    if count>limit:
        break;
    for j in range(W,height-W):
        if count>limit:
            break;
        count=count+1
        k=1
        dic={}
        for x in range(i-W, i+W+1):
            for y in range(j-W,j+W+1):
                if x==i and y==j:
                    continue
                dic["R"+str(k)] = pix[x, y][0]
                dic["G"+str(k)] = pix[x, y][1]
                dic["B"+str(k)] = pix[x, y][2]
                k=k+1     
        ls_defects.append(dic)    
data_test_defects=pd.DataFrame(ls_defects)
predicted_defects = model_defect.predict(data_test_defects)
print("Построена модель дефектов")
#print(predicted_defects)

#теперь будем восстанавливать картинку
index=0
for i in range(W,width-W):
    if count>limit:
        break;
    for j in range(W,height-W):
        if count>limit:
            break;
        count=count+1        
        dic={}
        if predicted_defects[index]>0.5:
            draw.point((i, j), (pix[i, j][0], pix[i, j][1], pix[i, j][2])) 
            #draw.point((i, j), (255, 255, 255))
        else:
            #draw.point((i, j), (0, 0, 0))
            k=1
            for x in range(i-W, i+W+1):
                for y in range(j-W,j+W+1):
                    if x==i and y==j:
                        continue
                    dic["R"+str(k)] = pix[x, y][0]
                    dic["G"+str(k)] = pix[x, y][1]
                    dic["B"+str(k)] = pix[x, y][2]
                    k=k+1     
            dic["RY"]=pix[i, j][0]
            dic["GY"]=pix[i, j][1]
            dic["BY"]=pix[i, j][2]                
            ls_img=[dic]
            data_img=pd.DataFrame(ls_img)     
            #print(data_img)
            r = GetColor(data_img,["GY","BY"],modelR) #узнаём значение красного цвета пикселя
            g = GetColor(data_img,["RY","BY"],modelG) #зелёного
            b = GetColor(data_img,["RY","GY"],modelB) #синего        
            draw.point((i, j), (r, g, b)) #рисуем пиксель
        index=index+1
            
new_image.save("d:\\rep\\restored.png", "PNG") #не забываем сохранить изображение
print("Преобразование картинки успешно выполнено")

Теперь смотрим результат работы данной программы:

Как видим, ничего хорошего не получилось. Наоборот, программа сделала еще хуже. Давайте разбираться, почему. На прошлом уроке было сказано, что треть дефектов мы не исправим, и примерно 20% пикселей будут заменены на предсказанные, но только с другой картинки. По сути, так и вышло.

Давайте попробуем заменить дефектные пиксели чёрной точкой, а нормальные белой. Тогда у нас получиться вот что:

Как видим, качество поиска дефектных пикселей очень низкое. Что тут можно сделать? Есть три пути:

  1. Попробовать обучить на другой картинке, более близкой к исходной.
  2. Попробовать улучить модель. Например, вместо случайного леса взять другой алгоритм
  3. По другому сформировать обучающую выборку. Например, загонять в модель другое количество входных точек, или вместо точек использовать что-то другое, например, какие-либо текстурные признаки.

Давайте, для начала, посмотрим, а как этот алгоритм сработает на той картинке, где он обучался:

 

Как видим, и тут дефекты ищутся некачественно, хотя и лучше, чем на другой картинке. Вангую, картинка тоже не восстановиться, если ее прогнать через алгоритм. Смотрим:

 

Собственно, так и получилось. Так что первый вариант попробовать другую картинку отпадает, так как даже на самой себе не получается. Осталось еще две идеи: взять другую модель и по другому сформировать выборку. Но это мы сделаем на следующих уроках.

Следующий урок.

Comments

So empty here ... leave a comment!

Добавить комментарий

Sidebar