修改文件名属性

generator.py

@@ -118,7 +118,9 @@ class DataMgr:
             self.pre_file = None
 
     def encode(self, cp_points: defaultdict[str], cp_nozzle: defaultdict[str], board_width, board_height):
-
+        assert len(cp_points.keys()) == len(cp_nozzle.keys())
+        assert len(cp_nozzle.keys()) <= self.max_component_types and len(
+            set(cp_nozzle.values())) <= self.max_nozzle_types
         # === general info ===
         total_points = sum(points for points in cp_points.values())
         total_component_types, total_nozzle_types = len(cp_points.keys()), len(set(cp_nozzle.values()))
@@ -127,7 +129,7 @@ class DataMgr:
         data.extend([board_width, board_height])
 
         # === heuristic info ===
-        cycle, nozzle_change, anc_move, pickup = self.heuristic_estimator(cp_points, cp_nozzle)
+        cycle, nozzle_change, anc_move, pickup = self.heuristic_objective(cp_points, cp_nozzle)
         data.extend([cycle, nozzle_change, anc_move, pickup])
 
         # === nozzle info ===
@@ -159,7 +161,7 @@ class DataMgr:
         # for _ in range(self.max_component_types - total_component_types):
         #     data.extend([0 for _ in range(self.max_nozzle_types)])
 
-        # === new component info ===
+        # === component info ===
         comp_data_slice = defaultdict(list)
         for idx in range(self.max_nozzle_types):
             comp_data_slice[idx] = []
@@ -180,7 +182,10 @@ class DataMgr:
             data.extend(comp_data_slice[idx])
 
         return data
-    def heuristic_estimator(self, cp_points, cp_nozzle):
+
+    def heuristic_objective(self, cp_points, cp_nozzle):
+        if len(cp_points.keys()) == 0:
+            return 0
         nozzle_heads, nozzle_points = defaultdict(int), defaultdict(int)
         for idx, points in cp_points.items():
             if points == 0:
@@ -400,8 +405,9 @@ class DataMgr:
         return pcb_data, component_data
 
     def loader(self, file_path):
-        train_data, time_data = [], []
-        cycle_data, nozzle_change_data, anc_move_data, pickup_data, movement_data, point_data = [], [], [], [], [], []
+        input_data, output_data = [], []        # 输入数据包含元件点数、吸嘴信息等，输出信息包含组装时间
+        # cycle_data, nozzle_change_data, anc_move_data, pickup_data, point_data = [], [], [], [], []
+        # pick_move_data, place_move_data = [], []
         with open(file_path, 'r') as file:
             line = file.readline()
             while line:
@@ -416,33 +422,56 @@ class DataMgr:
 
                     component_type, nozzle_type = items[12 + cp_idx * 3], items[13 + cp_idx * 3]
                     cp_points[component_type], cp_nozzle[component_type] = points, nozzle_type
+
                     # cp_width[component_type], cp_height[component_type] = float(items[15 + cp_idx * 5]), float(
                     #     items[16 + cp_idx * 5])
 
-                # if len(set(cp_nozzle.values())) > 2 or len(set(cp_nozzle.keys())) > 3:
-                if len(cp_points.keys()) > 30:
+                if len(cp_points.keys()) > 20 or len(cp_points.keys()) < 5:
                     line = file.readline()
                     continue
 
-                cycle_data.append(float(items[1]))
-                nozzle_change_data.append(float(items[2]))
-                anc_move_data.append(float(items[3]))
-                pickup_data.append(float(items[4]))
-                movement_data.append(float(items[5]) + float(items[6]))
-                point_data.append(sum(pt for pt in cp_points.values()))
+                board_width, board_height = float(items[7]), float(items[8])
+
+                # cycle_data.append(float(items[1]))
+                # nozzle_change_data.append(float(items[2]))
+                # anc_move_data.append(float(items[3]))
+                # pickup_data.append(float(items[4]))
+                # pick_move_data.append(float(items[5]))
+                # place_move_data.append(float(items[6]))
+                # point_data.append(sum(pt for pt in cp_points.values()))
 
                 # assembly time data
-                time_data.append(float(items[0]))
+                output_data.append(float(items[0]))
 
-                train_data.append(self.encode(cp_points, cp_nozzle, float(items[7]), float(items[8])))
+                # train_data.append(self.encode(cp_points, cp_nozzle, float(items[7]), float(items[8])))
+                input_data.append([cp_points, cp_nozzle, board_width, board_height])
                 # train_data[-1].extend([cycle_data[-1], nozzle_change_data[-1], anc_move_data[-1], pickup_data[-1]])
                 line = file.readline()
 
-        return train_data, time_data, cycle_data, nozzle_change_data, anc_move_data, pickup_data, point_data
+        # return train_data, time_data, cycle_data, nozzle_change_data, anc_move_data, pickup_data, pick_move_data, \
+        #        place_move_data, point_data
+
+        return [input_data, output_data]
+
+    def neural_encode(self, input_data):
+        train_data = []
+        for cp_points, cp_nozzle, board_width, board_height in input_data:
+            train_data.append(self.encode(cp_points, cp_nozzle, board_width, board_height))
+        return train_data
 
     def get_feature(self):
         return (self.max_component_types + 2) * self.max_nozzle_types + 5 + 4
 
+    # def neural_encode(self, input_data):
+    #     train_data = []
+    #     for cp_points, cp_nozzle, board_width, board_height in input_data:
+    #         train_data.append(
+    #             [len(cp_points.keys()), len(cp_nozzle.keys()), sum(cp_points.values()), board_width, board_height])
+    #     return train_data
+    #
+    # def get_feature(self):
+    #     return 5
+
     def get_update_round(self):
         return self.update