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
https://i.imgur.com/PPaZHiu.png