train
@@ -44,7 +44,6 @@ func (a *App) Startup(ctx context.Context) {
|
|||||||
a.ctx = ctx
|
a.ctx = ctx
|
||||||
|
|
||||||
// 设置 ONNX Runtime DLL 路径
|
// 设置 ONNX Runtime DLL 路径
|
||||||
// ort.SetSharedLibraryPath("onnxruntime.dll")
|
|
||||||
exeDir, _ := os.Executable()
|
exeDir, _ := os.Executable()
|
||||||
println("exeDir", exeDir)
|
println("exeDir", exeDir)
|
||||||
ort.SetSharedLibraryPath(filepath.Join(filepath.Dir(exeDir), "onnxruntime.dll"))
|
ort.SetSharedLibraryPath(filepath.Join(filepath.Dir(exeDir), "onnxruntime.dll"))
|
||||||
@@ -160,7 +159,6 @@ func (a *App) GormDB() (*gorm.DB, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (a *App) UploadImage(data []byte, filename string) Response {
|
func (a *App) UploadImage(data []byte, filename string) Response {
|
||||||
println("UploadImage")
|
|
||||||
uploadsDir := publicImagePath
|
uploadsDir := publicImagePath
|
||||||
err := os.MkdirAll(uploadsDir, 0755)
|
err := os.MkdirAll(uploadsDir, 0755)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
@@ -26,22 +26,25 @@ if mode == "toy":
|
|||||||
]
|
]
|
||||||
num_classes = len(label_name)
|
num_classes = len(label_name)
|
||||||
elif mode == "benchmark":
|
elif mode == "benchmark":
|
||||||
epoch = 100 # convnext->100, resnet18->50
|
epoch = 50 # resnet=80, resnet18=50, convnext=100
|
||||||
lr = 1e-4
|
lr = 1e-4 # resnet, resnet18=1e-4
|
||||||
|
weight_decay = 1e-4 # resnet=1e-3, resnet18=1e-4
|
||||||
batch_size = 8
|
batch_size = 8
|
||||||
input_size = 224
|
input_size = 224
|
||||||
|
|
||||||
# 分类
|
# 分类
|
||||||
label_name = [
|
label_name = [
|
||||||
"american_shorthair",
|
"american_shorthair", # 美国短毛猫
|
||||||
"bengal",
|
"british_shorthair", # 英国短毛猫
|
||||||
"british_shorthair",
|
"ragdoll", # 布偶猫
|
||||||
"exotic_shorthair",
|
"exotic_shorthair", # 异国短毛猫
|
||||||
"maine_coon",
|
"maine_coon", # 缅因猫
|
||||||
"ragdoll",
|
"siamese", # 暹罗猫
|
||||||
"scottish_fold",
|
"sphynx", # 斯芬克斯猫
|
||||||
"siamese",
|
"turkish_van", # 土耳其梵猫
|
||||||
"sphynx",
|
"bengal", # 孟加拉豹猫
|
||||||
"turkish_van",
|
"scottish_fold", # 苏格兰折耳猫
|
||||||
|
"none", # 风景人物
|
||||||
|
"other", # 其他动物
|
||||||
]
|
]
|
||||||
num_classes = len(label_name)
|
num_classes = len(label_name)
|
||||||
@@ -77,3 +77,4 @@ print("test_dataset", len(test_dataset))
|
|||||||
|
|
||||||
train_dataloader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
|
train_dataloader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
|
||||||
test_dataloader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
|
test_dataloader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
|
||||||
|
|
||||||
|
|||||||
@@ -1,67 +0,0 @@
|
|||||||
class,image_count,avg_width,avg_height,min_width,min_height,max_width,max_height,formats,corrupt_files
|
|
||||||
abyssinian,200,148,117,88,46,162,140,"jpeg, png",0
|
|
||||||
cyprus,200,159,100,93,40,162,140,"jpeg, png",0
|
|
||||||
lykoi,200,142,122,68,55,162,140,"jpeg, png",0
|
|
||||||
donskoy,200,141,123,67,54,162,140,jpeg,0
|
|
||||||
chausie,200,143,119,78,46,162,140,"jpeg, png",0
|
|
||||||
european_shorthair,200,149,118,70,49,162,140,"jpeg, png",0
|
|
||||||
turkish_van,200,149,120,87,56,162,140,"jpeg, png",0
|
|
||||||
pixie_bob,200,147,119,78,49,162,140,"jpeg, png",0
|
|
||||||
ragdoll,199,149,116,78,50,162,140,jpeg,0
|
|
||||||
german_rex,199,138,125,72,50,162,140,"jpeg, png",0
|
|
||||||
american_shorthair,199,147,114,65,50,162,140,"jpeg, png",0
|
|
||||||
sokoke,199,145,122,70,49,162,140,"jpeg, png",0
|
|
||||||
khao_manee,198,144,120,78,49,162,140,"jpeg, png",0
|
|
||||||
thai,198,143,118,68,60,162,140,"jpeg, png",0
|
|
||||||
cymric,198,148,120,63,43,162,140,"jpeg, png",0
|
|
||||||
oriental_shorthair,197,141,123,78,54,162,140,jpeg,0
|
|
||||||
cornish_rex,197,147,117,75,56,162,140,"jpeg, png",0
|
|
||||||
burmese,197,147,117,45,58,162,140,"jpeg, png",0
|
|
||||||
savannah,197,153,106,78,40,162,140,"jpeg, png",0
|
|
||||||
american_wirehair,196,149,117,76,50,162,140,"jpeg, png",0
|
|
||||||
peterbald,196,144,123,77,50,162,140,"jpeg, png",0
|
|
||||||
karelian_bobtail,196,145,125,67,65,162,140,"jpeg, png",0
|
|
||||||
tonkinese,195,148,116,70,53,162,140,"jpeg, png",0
|
|
||||||
balinese,195,154,112,79,54,162,140,jpeg,0
|
|
||||||
japanese_bobtail,194,149,118,87,49,162,140,"jpeg, png",0
|
|
||||||
nebelung,194,143,119,64,46,162,140,jpeg,0
|
|
||||||
selkirk_rex,192,144,122,78,76,162,140,jpeg,0
|
|
||||||
persian,192,150,114,78,46,162,140,jpeg,0
|
|
||||||
manx,192,155,113,86,50,162,140,"jpeg, png",0
|
|
||||||
himalayan,192,158,109,93,40,162,140,jpeg,0
|
|
||||||
munchkin,191,147,117,61,48,162,140,jpeg,0
|
|
||||||
bengal,189,151,115,78,71,162,140,jpeg,0
|
|
||||||
turkish_angora,188,145,119,66,52,162,140,jpeg,0
|
|
||||||
vankedisi,187,147,116,63,47,162,140,"jpeg, png",0
|
|
||||||
scottish_fold,184,143,120,78,50,162,140,jpeg,0
|
|
||||||
egyptian_mau,184,144,121,72,74,162,140,"jpeg, png",0
|
|
||||||
ocicat,182,150,115,62,44,162,140,"jpeg, png",0
|
|
||||||
ragamuffin,182,149,116,78,44,162,140,"jpeg, png",0
|
|
||||||
serengeti,175,159,100,78,53,300,140,"jpeg, png",0
|
|
||||||
british_shorthair,174,148,117,78,50,162,140,jpeg,0
|
|
||||||
toyger,160,145,120,78,50,162,140,"jpeg, png",0
|
|
||||||
siberian,159,153,116,78,55,162,140,jpeg,0
|
|
||||||
havana_brown,159,133,127,75,34,300,140,"jpeg, png",0
|
|
||||||
exotic_shorthair,157,148,118,93,55,300,140,"jpeg, png",0
|
|
||||||
bombay,154,139,123,46,50,162,140,"jpeg, png",0
|
|
||||||
korat,152,147,117,93,50,162,140,"jpeg, png",0
|
|
||||||
safari,150,158,105,79,38,300,140,"jpeg, png",0
|
|
||||||
american_bobtail,140,155,114,64,48,162,140,jpeg,0
|
|
||||||
mekong_bobtail,140,149,118,44,50,162,140,"jpeg, png",0
|
|
||||||
korean_bobtail,139,140,129,63,81,162,140,jpeg,0
|
|
||||||
siamese,139,152,115,78,53,300,140,"jpeg, png",0
|
|
||||||
somali,139,150,112,61,53,162,140,jpeg,0
|
|
||||||
devon_rex,138,150,116,78,55,162,140,jpeg,0
|
|
||||||
american_curl,138,155,112,91,51,300,140,"jpeg, png",0
|
|
||||||
ural_rex,137,144,121,48,65,162,140,"jpeg, png",0
|
|
||||||
singapura,136,158,109,93,54,300,140,"jpeg, png",0
|
|
||||||
ukrainian_levkoy,134,140,123,78,59,162,140,jpeg,0
|
|
||||||
maine_coon,133,141,122,66,79,162,140,jpeg,0
|
|
||||||
birman,131,155,113,93,56,162,140,jpeg,0
|
|
||||||
oregon_rex,121,147,123,85,66,300,140,"jpeg, png",0
|
|
||||||
kurilian_bobtail,120,149,121,78,81,162,140,jpeg,0
|
|
||||||
laperm,120,149,116,78,55,162,140,"jpeg, png",0
|
|
||||||
sphynx,120,149,117,47,63,162,140,jpeg,0
|
|
||||||
chartreux,114,146,117,75,61,162,140,jpeg,0
|
|
||||||
russian_blue,109,152,111,88,36,162,140,jpeg,0
|
|
||||||
norwegian_forest_cat,97,150,114,78,56,162,140,jpeg,0
|
|
||||||
|
@@ -4,13 +4,12 @@ import glob
|
|||||||
import torch
|
import torch
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
import onnxruntime as ort # 【修改1】替换torch导入
|
import onnxruntime as ort
|
||||||
from torchvision import transforms # 保留transforms,仍用于预处理
|
from torchvision import transforms
|
||||||
from core.const import mode, label_name, input_size
|
from core.const import mode, label_name, input_size
|
||||||
|
|
||||||
|
|
||||||
def test():
|
def test():
|
||||||
# 【修改2】选择ONNX Runtime的执行提供程序,自动选择CPU或CUDA
|
|
||||||
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if 'CUDAExecutionProvider' in ort.get_available_providers() else ['CPUExecutionProvider']
|
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if 'CUDAExecutionProvider' in ort.get_available_providers() else ['CPUExecutionProvider']
|
||||||
print(f"使用设备: {providers[0]}")
|
print(f"使用设备: {providers[0]}")
|
||||||
|
|
||||||
@@ -19,7 +18,7 @@ def test():
|
|||||||
model_dir = os.path.join(core_dir, "models")
|
model_dir = os.path.join(core_dir, "models")
|
||||||
dataset_dir = os.path.join(core_dir, "dataset", mode, "test")
|
dataset_dir = os.path.join(core_dir, "dataset", mode, "test")
|
||||||
|
|
||||||
# 【修改3】加载ONNX模型,替代原来的PyTorch模型加载
|
# 加载ONNX模型
|
||||||
session = ort.InferenceSession(os.path.join(model_dir, "resnet_epoch_100.onnx"), providers=providers)
|
session = ort.InferenceSession(os.path.join(model_dir, "resnet_epoch_100.onnx"), providers=providers)
|
||||||
|
|
||||||
# 获取输入名称(用于后续推理时指定输入)
|
# 获取输入名称(用于后续推理时指定输入)
|
||||||
@@ -28,7 +27,7 @@ def test():
|
|||||||
im_list = glob.glob(os.path.join(dataset_dir, "*", "*.jpg"))
|
im_list = glob.glob(os.path.join(dataset_dir, "*", "*.jpg"))
|
||||||
np.random.shuffle(im_list)
|
np.random.shuffle(im_list)
|
||||||
|
|
||||||
# 预处理完全不变
|
# 预处理
|
||||||
test_transform = transforms.Compose([
|
test_transform = transforms.Compose([
|
||||||
transforms.Resize(input_size),
|
transforms.Resize(input_size),
|
||||||
transforms.CenterCrop(input_size),
|
transforms.CenterCrop(input_size),
|
||||||
@@ -42,15 +41,15 @@ def test():
|
|||||||
inputs = test_transform(im_data)
|
inputs = test_transform(im_data)
|
||||||
inputs = torch.unsqueeze(inputs, dim=0) # 这里还在用torch,下面会改
|
inputs = torch.unsqueeze(inputs, dim=0) # 这里还在用torch,下面会改
|
||||||
|
|
||||||
# 【修改4】将输入转为numpy,ONNX Runtime需要numpy输入
|
# 将输入转为numpy,ONNX Runtime需要numpy输入
|
||||||
inputs = inputs.numpy()
|
inputs = inputs.numpy()
|
||||||
if providers[0] == 'CPUExecutionProvider':
|
if providers[0] == 'CPUExecutionProvider':
|
||||||
inputs = inputs.astype(np.float32)
|
inputs = inputs.astype(np.float32)
|
||||||
|
|
||||||
# 【修改5】ONNX Runtime推理,输出直接是numpy数组
|
# ONNX Runtime推理,输出直接是numpy数组
|
||||||
outputs = session.run(None, {input_name: inputs})[0]
|
outputs = session.run(None, {input_name: inputs})[0]
|
||||||
|
|
||||||
# 【修改6】解析结果,直接用numpy操作
|
# 解析结果,直接用numpy操作
|
||||||
pred = np.argmax(outputs, axis=1)
|
pred = np.argmax(outputs, axis=1)
|
||||||
print(label_name[pred[0]], " ", im_path)
|
print(label_name[pred[0]], " ", im_path)
|
||||||
|
|
||||||
|
|||||||
@@ -21,7 +21,7 @@ def test():
|
|||||||
print("model_dir", model_dir)
|
print("model_dir", model_dir)
|
||||||
|
|
||||||
net = resnet()
|
net = resnet()
|
||||||
net.load_state_dict(torch.load(os.path.join(model_dir, "resnet_epoch_99.pth"), weights_only=True))
|
net.load_state_dict(torch.load(os.path.join(model_dir, "resnet_epoch_80.pth"), weights_only=True))
|
||||||
|
|
||||||
print("111")
|
print("111")
|
||||||
|
|
||||||
|
|||||||
@@ -18,10 +18,8 @@ def test():
|
|||||||
model_dir = os.path.join(core_dir, "models")
|
model_dir = os.path.join(core_dir, "models")
|
||||||
dataset_dir = os.path.join(core_dir, "dataset", mode, "test")
|
dataset_dir = os.path.join(core_dir, "dataset", mode, "test")
|
||||||
|
|
||||||
print("model_dir", model_dir)
|
|
||||||
|
|
||||||
net = resnet18()
|
net = resnet18()
|
||||||
net.load_state_dict(torch.load(os.path.join(model_dir, "resnet18_epoch_50_bak2.pth"), weights_only=True))
|
net.load_state_dict(torch.load(os.path.join(model_dir, "resnet18_epoch_50.pth"), weights_only=True))
|
||||||
|
|
||||||
im_list = glob.glob(os.path.join(dataset_dir, "*", "*.jpg"))
|
im_list = glob.glob(os.path.join(dataset_dir, "*", "*.jpg"))
|
||||||
np.random.shuffle(im_list)
|
np.random.shuffle(im_list)
|
||||||
@@ -35,6 +33,8 @@ def test():
|
|||||||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
||||||
])
|
])
|
||||||
|
|
||||||
|
ary = []
|
||||||
|
errors = []
|
||||||
for im_path in im_list:
|
for im_path in im_list:
|
||||||
net.eval()
|
net.eval()
|
||||||
im_data = Image.open(im_path)
|
im_data = Image.open(im_path)
|
||||||
@@ -46,13 +46,25 @@ def test():
|
|||||||
outputs = net.forward(inputs)
|
outputs = net.forward(inputs)
|
||||||
|
|
||||||
_, pred = torch.max(outputs.data, dim=1)
|
_, pred = torch.max(outputs.data, dim=1)
|
||||||
print(label_name[pred.cpu().numpy()[0]], " ", im_path)
|
# print(label_name[pred.cpu().numpy()[0]], " ", im_path)
|
||||||
|
|
||||||
img = np.asarray(im_data)
|
# img = np.asarray(im_data)
|
||||||
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
# img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
||||||
img = cv2.resize(img, (200, 200))
|
# img = cv2.resize(img, (200, 200))
|
||||||
cv2.imshow("img", img)
|
# cv2.imshow("img", img)
|
||||||
cv2.waitKey(0)
|
# cv2.waitKey(0)
|
||||||
|
result = label_name[pred.cpu().numpy()[0]]
|
||||||
|
|
||||||
|
if result in im_path:
|
||||||
|
ary.append(True)
|
||||||
|
else:
|
||||||
|
ary.append(False)
|
||||||
|
print(f"{label_name[pred.cpu().numpy()[0]]} {im_path}\n")
|
||||||
|
errors.append(f"{label_name[pred.cpu().numpy()[0]]} {im_path}")
|
||||||
|
|
||||||
|
print("ary", ary)
|
||||||
|
# print("errors", errors)
|
||||||
|
print(sum(ary) / len(ary))
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
@@ -2,7 +2,7 @@ import os
|
|||||||
import torch
|
import torch
|
||||||
from core.nets.resnet import resnet
|
from core.nets.resnet import resnet
|
||||||
from core.dataloader.dataloader import train_dataloader
|
from core.dataloader.dataloader import train_dataloader
|
||||||
from core.const import epoch, lr, batch_size
|
from core.const import epoch, lr, weight_decay, batch_size
|
||||||
|
|
||||||
|
|
||||||
def train():
|
def train():
|
||||||
@@ -13,9 +13,9 @@ def train():
|
|||||||
|
|
||||||
loss_func = torch.nn.CrossEntropyLoss()
|
loss_func = torch.nn.CrossEntropyLoss()
|
||||||
|
|
||||||
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
|
optimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
|
||||||
|
|
||||||
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
|
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
|
||||||
|
|
||||||
for e in range(epoch):
|
for e in range(epoch):
|
||||||
print("epoch: ", e)
|
print("epoch: ", e)
|
||||||
|
|||||||
@@ -2,7 +2,7 @@ import os
|
|||||||
import torch
|
import torch
|
||||||
from core.nets.resnet18 import resnet18
|
from core.nets.resnet18 import resnet18
|
||||||
from core.dataloader.dataloader import train_dataloader
|
from core.dataloader.dataloader import train_dataloader
|
||||||
from core.const import epoch, lr, batch_size
|
from core.const import epoch, lr, weight_decay, batch_size
|
||||||
|
|
||||||
|
|
||||||
def train():
|
def train():
|
||||||
@@ -13,7 +13,7 @@ def train():
|
|||||||
|
|
||||||
loss_func = torch.nn.CrossEntropyLoss()
|
loss_func = torch.nn.CrossEntropyLoss()
|
||||||
|
|
||||||
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
|
optimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
|
||||||
|
|
||||||
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20, eta_min=1e-6) # 余弦退火
|
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20, eta_min=1e-6) # 余弦退火
|
||||||
|
|
||||||
|
|||||||
|
After Width: | Height: | Size: 33 KiB |
|
After Width: | Height: | Size: 64 KiB |
|
After Width: | Height: | Size: 26 KiB |
|
After Width: | Height: | Size: 24 KiB |
|
After Width: | Height: | Size: 34 KiB |
|
After Width: | Height: | Size: 54 KiB |
|
After Width: | Height: | Size: 16 KiB |
|
After Width: | Height: | Size: 29 KiB |
|
After Width: | Height: | Size: 26 KiB |
|
After Width: | Height: | Size: 27 KiB |
|
After Width: | Height: | Size: 16 KiB |
|
After Width: | Height: | Size: 31 KiB |