Files
smt-optimizer/dataloader.py

117 lines
5.2 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

from base_optimizer.optimizer_common import *
def load_data(filename: str, default_feeder_limit=1, load_cp_data=True, load_feeder_data=True, cp_auto_register=False):
# 读取PCB数据
filename = 'data/' + filename
part_content, step_content = False, False
part_start_line, step_start_line, part_end_line, step_end_line = -1, -1, -1, -1
line_counter = 0
with open(filename, 'r') as file:
line = file.readline()
while line:
if line == '[Part]\n':
part_content = True
part_start_line = line_counter
elif line == '[Step]\n':
step_content = True
step_start_line = line_counter
elif line == '\n':
if part_content:
part_content = False
part_end_line = line_counter
elif step_content:
step_content = False
step_end_line = line_counter
line_counter += 1
line = file.readline()
if part_content:
part_end_line = line_counter
elif step_content:
step_end_line = line_counter
pcb_data = pd.DataFrame(
pd.read_csv(filepath_or_buffer=filename, skiprows=step_start_line + 1, nrows=step_end_line - step_start_line + 1,
sep='\t', header=None))
if len(pcb_data.columns) <= 17:
step_col = ["ref", "x", "y", "z", "r", "part", "desc", "fdr", "nz", "hd", "cs", "cy", "sk", "bl", "ar",
"pl", "lv"]
elif len(pcb_data.columns) <= 18:
step_col = ["ref", "x", "y", "z", "r", "part", "desc", "fdr", "nz", "hd", "cs", "cy", "sk", "bl", "ar", "fid",
"pl", "lv"]
else:
step_col = ["ref", "x", "y", "z", "r", "part", "desc", "fdr", "nz", "hd", "cs", "cy", "sk", "bl", "ar", "fid",
"", "pl", "lv"]
pcb_data.columns = step_col
pcb_data = pcb_data.dropna(axis=1)
# 坐标系处理
# pcb_data = pcb_data.sort_values(by = ['x', 'y'], ascending = True)
# pcb_data["x"] = pcb_data["x"].apply(lambda x: -100+x)
# 注册元件检查
part_feeder_assign = defaultdict(set)
part_col = ["part", "fdr", "nz", 'feeder-limit']
try:
if part_start_line != -1:
component_data = pd.DataFrame(
pd.read_csv(filepath_or_buffer=filename, sep='\t', header=None, skiprows=part_start_line + 1,
nrows=part_end_line - part_start_line - 1))
component_data.columns = part_col
else:
component_data = pd.DataFrame(columns=part_col)
except:
component_data = pd.DataFrame(columns=part_col)
component_data['points'] = 0
part_col.append('points')
for _, data in pcb_data.iterrows():
part, nozzle = data.part, data.nz.split(' ')[1]
slot = data['fdr'].split(' ')[0]
if part not in component_data['part'].values:
if not cp_auto_register:
raise Exception("unregistered component: " + component_data['part'].values)
else:
component_data = pd.concat([component_data, pd.DataFrame(
[part, '', 'SM8', nozzle, default_feeder_limit, 0], index=part_col).T],
ignore_index=True)
warning_info = 'register component ' + part + ' with default feeder type'
warnings.warn(warning_info, UserWarning)
part_index = component_data[component_data['part'] == part].index.tolist()[0]
part_feeder_assign[part].add(slot)
component_data.loc[part_index, 'points'] += 1
if nozzle != 'A' and component_data.loc[part_index, 'nz'] != nozzle:
warning_info = 'the nozzle type of component ' + part + ' is not consistent with the pcb data'
warnings.warn(warning_info, UserWarning)
# 清除点数为0的数据
component_data = component_data[component_data['points'] != 0].reset_index(drop=True)
for idx, data in component_data.iterrows():
if data['fdr'][0:3] == 'SME': # 电动供料器和气动供料器参数一致
component_data.at[idx, 'fdr'] = data['fdr'][0:2] + data['fdr'][3:]
# 读取供料器基座数据
feeder_data = pd.DataFrame(columns=['slot', 'part', 'arg']) # arg表示是否为预分配不表示分配数目
if load_feeder_data:
for _, data in pcb_data.iterrows():
slot, part = data['fdr'].split(' ')
if slot[0] != 'F' and slot[0] != 'R':
continue
slot = int(slot[1:]) if slot[0] == 'F' else int(slot[1:]) + max_slot_index // 2
feeder_data = pd.concat([feeder_data, pd.DataFrame([slot, part, 1]).T])
feeder_data.drop_duplicates(subset='slot', inplace=True, ignore_index=True)
# 随机移除部分已安装的供料器
if load_feeder_data == 2:
drop_index = random.sample(list(range(len(feeder_data))), len(feeder_data) // 2)
feeder_data.drop(index=drop_index, inplace=True)
feeder_data.sort_values(by='slot', ascending=True, inplace=True, ignore_index=True)
pcb_data = pcb_data.sort_values(by="x", ascending=False)
return pcb_data, component_data, feeder_data