import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib
import os
import cv2
from PIL import Image
import torchvision.transforms as transforms
import torch.optim as optim
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
from Test.CNN import Net
import json
from Test.train_data import Mydataset,pad_image

# 構(gòu)建神經(jīng)網(wǎng)絡(luò)
class Net(nn.Module):#定義網(wǎng)絡(luò)模塊
    def __init__(self):
        super(Net, self).__init__()
        # 卷積，該圖片有3層，6個(gè)特征，長(zhǎng)寬均為5*5的像素點(diǎn)，每隔1步跳一下
        self.conv1 = nn.Conv2d(3, 6, 5)
        #//(conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
        self.pool = nn.MaxPool2d(2, 2)#最大池化
        #//(pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        self.conv2 = nn.Conv2d(6, 16, 5)#卷積
        #//(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
        self.fc1 = nn.Linear(16*77*77, 120)#全連接層，圖片的維度為16，
        #(fc1): Linear(in_features=94864, out_features=120, bias=True)
        self.fc2 = nn.Linear(120, 84)#全連接層,輸入120個(gè)特征輸出84個(gè)特征
        self.fc3 = nn.Linear(84, 7)#全連接層，輸入84個(gè)特征輸出7個(gè)特征
 
   def forward(self, x):
        print("x.shape1: ", x.shape)
        x = self.pool(F.relu(self.conv1(x)))
        print("x.shape2: ", x.shape)
        x = self.pool(F.relu(self.conv2(x)))
        print("x.shape3: ", x.shape)
        x = x.view(-1, 16*77*77)
        print("x.shape4: ", x.shape)
        x = F.relu(self.fc1(x))
        print("x.shape5: ", x.shape)
        x = F.relu(self.fc2(x))
        print("x.shape6: ", x.shape)
        x = self.fc3(x)
        print("x.shape7: ", x.shape)
        return x

img_path = "TestData/test_data/1/Apple2 (1).jpg" #使用相對(duì)路徑
image = Image.open(img_path).convert('RGB')
image_pad = pad_image(image, (320, 320))
input = transform(image_pad).to(device).unsqueeze(0)
output = F.softmax(net(input), 1)
_, predicted = torch.max(output, 1)
score = float(output[0][predicted]*100)
print(class_map[predicted], " ", str(score)+" %")
plt.imshow(image_pad) # 顯示圖片