picutai.py

import torch
import torch.nn as nn
import cv2
import torch.optim as optim
import glob

#vec = (cv2.resize(cv2.cvtColor(cv2.imread('sand.jpg'),cv2.COLOR_BGR2GRAY),(28,28)) / 255.0)
#in_model = torch.from_numpy(vec).float().view(1,1,28,28)


# 2. 定义完整的模型（包含 Flatten 和 Linear）
model = nn.Sequential(
    nn.Conv2d(3, 8, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.Conv2d(8, 16, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(2), # 把上面的部分包进来
    nn.AdaptiveAvgPool2d((1, 1)),
    nn.Flatten(),      # 关键修复点
    nn.Linear(16, 2)
)

opt = optim.Adam(model.parameters(),lr=0.005)
CrossLoss = nn.CrossEntropyLoss()
for epoch in range(40):
    print("epoching...")
    for i in range(2):
        vec2 = (cv2.resize(cv2.imread(f'san{i}.jpg'),(28,28)) / 255.0)
        input_model = torch.from_numpy(vec2).float().view(1,3,28,28)

        if i == 0:
            label = 0
        else:
            label = 1

        tar = torch.LongTensor([label])
        output = model(input_model)
        print(output)
        loss = CrossLoss(output,tar)
        opt.zero_grad()
        loss.backward()
        opt.step()
        print(loss.item())
    
    #cv2.COLOR_BGR2GRAY
while True:
    model.eval()
    File = input("picuter File :")
    vec3 = (cv2.resize(cv2.imread(File),(28,28)) / 255.0)
    in_model = torch.from_numpy(vec3).float().view(1,3,28,28)
    with torch.no_grad():
        out = model(in_model)
        probab = torch.argmax(out,dim = 1).item() 
        print(probab)

        if probab > 0.5:
            print("This is dog.")
        else:
            print("Not dog!")
    kk = input("save model(Y/N):")
    if kk == "Y":
        torch.save(model,'model.pth')
    else:
        continue