mirror of
https://github.com/NanjingForestryUniversity/supermachine-tobacco.git
synced 2025-11-08 14:23:55 +00:00
54 lines
1.6 KiB
Python
54 lines
1.6 KiB
Python
#!/usr/bin/env python
|
|
# coding: utf-8
|
|
|
|
# # 模型的训练
|
|
|
|
# In[16]:
|
|
|
|
|
|
import numpy as np
|
|
import scipy
|
|
from imblearn.under_sampling import RandomUnderSampler
|
|
from models import AnonymousColorDetector
|
|
from utils import read_labeled_img
|
|
|
|
# ## 读取数据与构建数据集
|
|
|
|
# In[17]:
|
|
|
|
|
|
data_dir = "data/dataset"
|
|
color_dict = {(0, 0, 255): "yangeng", (255, 0, 0): 'beijing'}
|
|
label_index = {"yangeng": 1, "beijing": 0}
|
|
dataset = read_labeled_img(data_dir, color_dict=color_dict, is_ps_color_space=False)
|
|
rus = RandomUnderSampler(random_state=0)
|
|
x_list, y_list = np.concatenate([v for k, v in dataset.items()], axis=0).tolist(), \
|
|
np.concatenate([np.ones((v.shape[0],)) * label_index[k] for k, v in dataset.items()], axis=0).tolist()
|
|
|
|
x_resampled, y_resampled = rus.fit_resample(x_list, y_list)
|
|
dataset = {"inside": np.array(x_resampled)}
|
|
|
|
# ## 模型训练
|
|
|
|
# In[18]:
|
|
|
|
|
|
# 定义一些常量
|
|
threshold = 5
|
|
node_num = 20
|
|
negative_sample_num = None # None或者一个数字
|
|
world_boundary = np.array([0, 0, 0, 255, 255, 255])
|
|
# 对数据进行预处理
|
|
x = np.concatenate([v for k, v in dataset.items()], axis=0)
|
|
negative_sample_num = int(x.shape[0] * 1.2) if negative_sample_num is None else negative_sample_num
|
|
|
|
model = AnonymousColorDetector()
|
|
|
|
model.fit(x, world_boundary, threshold, negative_sample_size=negative_sample_num, train_size=0.7,
|
|
is_save_dataset=True, model_selection='dt')
|
|
# data = scipy.io.loadmat('dataset_2022-07-19_15-07.mat')
|
|
# x, y = data['x'], data['y'].ravel()
|
|
# model.fit(x, y=y, is_generate_negative=False, model_selection='dt')
|
|
|
|
model.save()
|