train.py

import math
from numpy.core.fromnumeric import shape
import torch
import matplotlib.pyplot as plt
import json
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms.functional as TF
import numpy as np
import cv2
import time
from PIL import Image
import random
import sys
import os
sys.path.append(os.path.abspath("C:/Users/x065p/Desktop/test/python/PyTorch/blender_vgg16/"))
from blender.gen_material_with_bsdf import blender_render
from vgg16 import Model

numClasses=27
mom_folder="C:/Users/x065p/Desktop/blender/result/image/python/"
losss=[]
epochs = 50
run_sample=1
# random.seed(4)
blender_object=0


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")



def bsdf_arg(arr):
    out =[]
    for i in arr:
        # print(type(i))
        try:
            for j in i:
                out.append(j)
        except:
            out.append(i)
    return out

    
f = open(mom_folder+"1000/save_value.txt",'r',encoding = 'utf-8')
datas=f.read()
datas=json.loads(datas)
model = Model().to(device)
optimizer = torch.optim.RMSprop(model.parameters(), lr=0.005)
loss_function = torch.nn.MSELoss()
blender_render_object=blender_render()


for epoch in range(epochs):
    input_img_arr=[]
    input_targe_arr=[]
    for j in range(run_sample):
        i=int(random.random()*1000)
        # print(i)
        # input_img = cv2.imread('data/1000/image/{}.png'.format(i))
        input_img = Image.open(mom_folder+'1000/{}.png'.format(i)).convert('RGB')
        input_img= input_img.resize((224,244))
        # input_img.show()
        input_img = TF.to_tensor(input_img)
        input_img=np.array(input_img)
        # print(input_img.shape)

        input_targe=bsdf_arg( datas[i])
        input_targe=np.array(input_targe)

        input_img_arr.append(input_img)
        input_targe_arr.append(input_targe)

    input_img_arr=torch.tensor(input_img_arr,requires_grad=True)
    input_img_arr=input_img_arr.float()
    input_img_arr=input_img_arr.to(device)

    input_targe_arr=torch.tensor(input_targe_arr)
    input_targe_arr=input_targe_arr.float()
    input_targe_arr=input_targe_arr.to(device)

    prediction = model(input_img_arr)
    predict_bsdf=prediction.detach().numpy()[0]

    blender_object=blender_render_object.img_by_bsdf(predict_bsdf,blender_object)
    # time.sleep(0.01)
    blender_image = Image.open(mom_folder+'tem{}.png'.format(j)).convert('RGB')
    blender_image= blender_image.resize((224,244))
    # blender_image.show()
    blender_image = TF.to_tensor(blender_image)
    blender_image = torch.tensor(blender_image,requires_grad=True)
    # blender_image=np.array(blender_image)
    # print(blender_image.shape)
    # print(input_img_arr[0].shape)
    loss = loss_function(blender_image, input_img_arr[0])
    losss.append(loss.item())

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    every=1
    if i % every == every-1:
        print("loss:",loss.item())
        print(epoch+1,"/",epochs,";",j+1,"/",run_sample)
        torch.save(model.state_dict(), mom_folder+"model_save.bin")
        # print("save")

with open(mom_folder+"loss.json","w") as f:
    f.write(json.dumps(losss))
    f.close()
    

#plt.plot(range(epochs), losss)
#plt.show()


vgg16.py
import torch
import torch.nn as nn
import torch.nn.functional as F

nn_net=[200*200*3,20*20*3,100,27]
# nn_net=[120000,27,100,27]
numClasses=27
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.vgg16 = nn.Sequential(
            nn.Conv2d(3, 64, 3),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, 3),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 256, 3),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 512, 3),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(inplace=True),
            # nn.Dropout(),
            nn.Linear(4096, 4096),
            # nn.ReLU(inplace=True),
            # nn.Dropout(),
            nn.Sigmoid(),
            nn.Linear(4096, numClasses),

        )
    def forward(self, x):
        x = self.vgg16(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x



log
C:\Users\x065p\Desktop\blender\blender file\python.blend\Text:94: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
loss: 0.07332869619131088
1
Saved: 'C:\Users\x065p\Desktop\blender\result\image\python\tem1.png'
 Time: 00:02.07 (Saving: 00:00.00)

loss: 0.024082180112600327
2
Saved: 'C:\Users\x065p\Desktop\blender\result\image\python\tem2.png'
 Time: 00:02.31 (Saving: 00:00.00)

loss: 0.08349122852087021
3
Saved: 'C:\Users\x065p\Desktop\blender\result\image\python\tem3.png'
 Time: 00:02.11 (Saving: 00:00.01)

loss: 0.07594556361436844
4
Saved: 'C:\Users\x065p\Desktop\blender\result\image\python\tem4.png'
 Time: 00:02.22 (Saving: 00:00.00)

loss: 0.09595263004302979
5
Saved: 'C:\Users\x065p\Desktop\blender\result\image\python\tem5.png'
 Time: 00:02.28 (Saving: 00:00.00)

#flow

1 model input: image


2 model predict  out :  27 length bsdf parameter
[0.45234,0.340502,0.9123...........]


3 blender render image by model predict bsdf parameter
                    
            
          



4 count loss with input image and blender render image