提交更新

script/trans_and_export_data.py

@@ -444,40 +444,77 @@ def shadow_casting(points,bbox):
   #print(points[flag])
   return flag
 
-def Save_Cloud_File(cloud):
+import numpy
+class NumpyEncoder(json.JSONEncoder):  
+    def default(self, obj):  
+        if isinstance(obj, (numpy.int_, numpy.intc, numpy.intp, numpy.int8,  
+            numpy.int16, numpy.int32, numpy.int64, numpy.uint8,  
+            numpy.uint16, numpy.uint32, numpy.uint64)):  
+            return int(obj)  
+        elif isinstance(obj, (numpy.float_, numpy.float16, numpy.float32,numpy.float64)):  
+            return float(obj)  
+        elif isinstance(obj, (numpy.ndarray,)):  
+            return obj.tolist()  
+        return json.JSONEncoder.default(self, obj) 
+
+
+g_saveFileType = 1 #0为pcd文件，1是bin文件
+g_showCloud = 0 #0不显示，1显示
+
+def Save_Cloud_File(cloud,json):
     #去除boxx内的点
     drawpointcloud = []
+    jsn_pcd = {}
+    jsn_pcd["fileuri"] = "5-3/pcd/" + cloud["pcdName"]
+    jsn_pcd["labels_box3D"] = []
     for bbox in cloud["boxes"]:
+        jsn_box = {}
+        jsn_box["box3D"] = {}
+        jsn_box["box3D"]["translation"] = {"x":bbox[4][0],"y":bbox[4][1],"z":bbox[4][2]}
+        jsn_box["box3D"]["size"] = {"x":bbox[6][0],"y":bbox[6][1],"z":bbox[6][2]}
+        jsn_box["box3D"]["rot_y"] = bbox[2]
+        jsn_box["category"] = bbox[0]
+        jsn_pcd["labels_box3D"].append(jsn_box)
+        
         flag_del = shadow_casting(cloud["np_pcd"],bbox[1])
         #g_np_flag = g_np_flag & (~flag_del)
         cloud["np_pcd"] = cloud["np_pcd"][~flag_del]
-        drawboxes = gen_o3d_3dbboxes(bbox[1])
-        drawpointcloud += [drawboxes]
-        heading_point = (bbox[3][0] + bbox[3][2]) / 2
-        heading_pcd = o3d.geometry.PointCloud()
-        heading_pcd.points = o3d.utility.Vector3dVector(np.array([heading_point]))
-        colors = [[0, 1, 0]]
-        heading_pcd.colors = o3d.utility.Vector3dVector(colors)
-        drawpointcloud += [heading_pcd]        
-    pcd_save = o3d.geometry.PointCloud()
-    pcd_save.points = o3d.utility.Vector3dVector(cloud["np_pcd"])
-    drawpointcloud += [pcd_save]
-    axis_pcd = o3d.geometry.TriangleMesh.create_coordinate_frame(size=5, origin=[0, 0, 0])
-    drawpointcloud += [axis_pcd]
-
-    savefile = os.path.join(cloud["path"],cloud["pcdName"])
-    o3d.io.write_point_cloud(savefile,pcd_save)
-
-    custom_draw_geometry_with_key_callback(drawpointcloud, savefile)
+        if g_showCloud == 1:
+            drawboxes = gen_o3d_3dbboxes(bbox[1])
+            drawpointcloud += [drawboxes]
+            heading_point = (bbox[3][0] + bbox[3][2]) / 2
+            heading_pcd = o3d.geometry.PointCloud()
+            heading_pcd.points = o3d.utility.Vector3dVector(np.array([heading_point]))
+            colors = [[0, 1, 0]]
+            heading_pcd.colors = o3d.utility.Vector3dVector(colors)
+            drawpointcloud += [heading_pcd]        
+    if g_saveFileType == 0 or g_showCloud == 1 :
+        pcd_save = o3d.geometry.PointCloud()
+        pcd_save.points = o3d.utility.Vector3dVector(cloud["np_pcd"])
+        if g_saveFileType == 0:
+            savefile = os.path.join(cloud["path"],cloud["pcdName"])
+            o3d.io.write_point_cloud(savefile,pcd_save)
+        if g_showCloud == 1:
+            drawpointcloud += [pcd_save]
+    if g_saveFileType == 1:
+        savefile = os.path.join(cloud["path"],cloud["pcdName"])
+        cloud["np_pcd"].tofile(savefile)
+    json["annotations"].append(jsn_pcd)
+    if g_showCloud == 1:
+        axis_pcd = o3d.geometry.TriangleMesh.create_coordinate_frame(size=5, origin=[0, 0, 0])
+        drawpointcloud += [axis_pcd]
+        custom_draw_geometry_with_key_callback(drawpointcloud, "")
+
+
     
 
-def Add_Cloud_box(save_cloud,add_idx,bbox,pickcloud):
+def Add_Cloud_box(save_cloud,add_idx,bbox,pickcloud,json):
     save_cloud[add_idx]["boxes"].append(bbox)
     delflag = in_hull(save_cloud[add_idx]["np_pcd"],bbox[1])
     save_cloud[add_idx]["np_pcd"] = save_cloud[add_idx]["np_pcd"][~delflag]
     save_cloud[add_idx]["np_pcd"] = np.vstack((save_cloud[add_idx]["np_pcd"],pickcloud))
     if save_cloud[0]["isSave"] == 1:
-        Save_Cloud_File(save_cloud[0])
+        Save_Cloud_File(save_cloud[0],json)
         save_cloud.pop(0)
 
 def Check_Add_Cloud_box(save_cloud,bbox,pcd,pcdName,path,index):
@@ -495,9 +532,9 @@ def Check_Add_Cloud_box(save_cloud,bbox,pcd,pcdName,path,index):
             break;
     if add_idx == -1 :
         #需要保存数据了
-        if index > 1 and len(save_cloud) >=  100 and  len(save_cloud[0]["boxes"]) >= 10  :
+        if index > 1 and len(save_cloud) >=  20 and  len(save_cloud[0]["boxes"]) >= 5  :
             save_cloud[0]["isSave"] = 1
-        elif index > 1 and len(save_cloud) > 0 and len(save_cloud[0]["boxes"]) > 50 :
+        elif index > 1 and len(save_cloud) > 0 and len(save_cloud[0]["boxes"]) > 10 :
             save_cloud[0]["isSave"] = 1
     if add_idx == -1:
         new_cloud = {}
@@ -554,6 +591,8 @@ if __name__ == '__main__':
     m =  [ 1.11119401e-03, -8.66741252e-04,  1.60358817e-01,  9.87057781e-01]
 
     save_cloud_list = [] #保存点云的数组
+    save_json = {}
+    save_json["annotations"] = []
 
     dir_pcd_list = {}
     dirs = os.listdir(origin_root_path)
@@ -612,8 +651,10 @@ if __name__ == '__main__':
                         continue;
                     #获取到汽车点云符合条件，可以添加到新点云里
                     # f1.write('%f,%f,%f,%f,%f,%f,%f,%f,%f,%f\n'%(exportCenterBL[0],exportCenterBL[1],bbox[4][0], bbox[4][1], bbox[4][2],bbox[6][0], bbox[6][1], bbox[6][2],angle2,laneAngle))
-                    
-                    idx = Check_Add_Cloud_box(save_cloud_list,bbox,g_converted_pcd,pcd_file.split("/")[-1],generate_pcd_path,index)
+                    name = pcd_file.split("/")[-1]
+                    if g_saveFileType == 1:
+                        name = name[0:-3] + "bin"
+                    idx = Check_Add_Cloud_box(save_cloud_list,bbox,g_converted_pcd,name,generate_pcd_path,index)
                     if idx == -1:
                         continue;
                     flag = in_hull(xyz,bbox[1])
@@ -628,13 +669,19 @@ if __name__ == '__main__':
                     print(pickcloud)
                     bbox[1] += dz
 
-                    Add_Cloud_box(save_cloud_list,idx,bbox,pickcloud)
+                    Add_Cloud_box(save_cloud_list,idx,bbox,pickcloud,save_json)
 
         if isDeal == 1:
             index += 1
+            if index >= 10:
+                break
         else:
             isStop = 1
-
+    for cloud in save_cloud_list:
+        Save_Cloud_File(cloud,save_json)
+    jsn_path = os.path.join(generate_root_path,generate_child_dir + ".json")
+    with open(jsn_path,'w') as file_obj:
+        json.dump(save_json,file_obj,cls=NumpyEncoder)
     # for dir in dirs:
     #     if os.path.isdir(origin_root_path + "/" + dir) == False:
     #         continue
@@ -703,6 +750,5 @@ if __name__ == '__main__':
     #             Add_Cloud_box(save_cloud_list,idx,bbox,pickcloud)
 
 
-    for cloud in save_cloud_list:
-        Save_Cloud_File(cloud)
+