69 lines
1.9 KiB
Python
69 lines
1.9 KiB
Python
import cv2
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from train_2 import train_model,prepare_data,predict
|
|
import joblib
|
|
|
|
inp = Path("data/train")
|
|
val = Path("data/val")
|
|
proc = Path("data/proc")
|
|
proc.mkdir(exist_ok=True)
|
|
|
|
|
|
def preproc(dir:Path):
|
|
d={}
|
|
for file in dir.glob("*/*.jpg"):
|
|
im = cv2.imread(file)
|
|
if im is None:
|
|
print(f"Error reading image: {file}")
|
|
continue
|
|
cl = file.parents[0].name
|
|
if cl not in d:
|
|
d[cl] = []
|
|
hsv = cv2.cvtColor(im,cv2.COLOR_BGR2HSV)
|
|
mask = hsv[:,:,1] > 150
|
|
# cor = np.argwhere(mask)
|
|
# y_min,x_min = cor.min(axis=0)
|
|
# y_max,x_max = cor.max(axis=0)
|
|
mask = mask[:,:,np.newaxis]
|
|
# cv2.findCoun
|
|
# im = im[y_min:y_max,x_min:x_max]
|
|
cnt = np.count_nonzero(mask)
|
|
hsv*=mask
|
|
h = round(np.sum(hsv[:,:,0])/cnt)
|
|
s = round(np.sum(hsv[:,:,1])/cnt)
|
|
v = round(np.sum(hsv[:,:,2])/cnt)
|
|
name = f"{h}_{s}_{v}.jpg"
|
|
d[cl].append((h,s,v))
|
|
# (proc/cl).mkdir(exist_ok=True)
|
|
# cv2.imwrite(proc/cl/name,cv2.cvtColor(hsv,cv2.COLOR_HSV2BGR))
|
|
return d
|
|
|
|
d:dict[str,list[tuple[int,int,int]]] = preproc(inp)
|
|
val:dict[str,list[tuple[int,int,int]]] = preproc(val)
|
|
|
|
print("数据预处理完成")
|
|
|
|
model, label_map = train_model(d)
|
|
print("训练完成")
|
|
joblib.dump(model, "model.pkl")
|
|
|
|
# model = joblib.load("model.pkl")
|
|
X_train, y_train = prepare_data(d)
|
|
print(predict(model, label_map, d))
|
|
print(predict(model, label_map, val))
|
|
# print(f"\n训练集准确率: {model.score(X_train, y_train):.5f}")
|
|
# X_train, y_train = prepare_data(val)
|
|
# print(f"\n训练集准确率: {model.score(X_train, y_train):.5f}")
|
|
# model.predict()
|
|
from src_predict import predictor
|
|
suc = cnt = 0
|
|
for file in Path("data/train").glob("*/*.jpg"):
|
|
cnt+=1
|
|
pcl,_=predictor(file)
|
|
acl = file.parents[0].name
|
|
if acl == pcl:
|
|
suc+=1
|
|
|
|
print(f"预测准确率: {suc/cnt:.4f}")
|