Merge branch 'test'

# Conflicts:
#	.gitattributes

.gitattributes

@@ -1,5 +1,5 @@
-vision_test/ merge=ours
-test/ merge=ours
+/vision_test/ merge=ours
+/test/ merge=ours
 
-.gitattributes merge=ours
-.gitignore merge=ours
+/.gitattributes merge=ours
+/.gitignore merge=ours

.gitignore

@@ -3,4 +3,3 @@
 **/build/
 **/install/
 **/log/
-**/VisionDetect/**/**.so

calibration_mat/eye_in_left_hand.json

@@ -1,28 +0,0 @@
-{
-    "T": [
-        [
-            -0.029312906184646676,
-            0.9941492398859674,
-            0.10396173510077301,
-            -0.0809046635564155
-        ],
-        [
-            -0.9993360095995784,
-            -0.02689542239474161,
-            -0.024579995358005674,
-            0.036195409684615196
-        ],
-        [
-            -0.021640088923136416,
-            -0.10461321660460321,
-            0.9942775173275502,
-            0.04913468358915477
-        ],
-        [
-            0.0,
-            0.0,
-            0.0,
-            1.0
-        ]
-    ]
-}

calibration_mat/eye_in_right_hand.json

@@ -1,28 +0,0 @@
-{
-    "T": [
-        [
-            0.02515190926712163,
-            0.9984434279780899,
-            -0.049780544267610596,
-            -0.08312977955463981
-        ],
-        [
-            -0.9996337211449028,
-            0.025617085950903107,
-            0.008728599966714646,
-            0.059044674332170574
-        ],
-        [
-            0.00999024575340213,
-            0.04954276975245833,
-            0.9987220378839358,
-            0.017378234075134728
-        ],
-        [
-            0.0,
-            0.0,
-            0.0,
-            1.0
-        ]
-    ]
-}

calibration_mat/eye_to_hand.json

@@ -1,28 +0,0 @@
-{
-    "T": [
-        [
-            0.01868075138315295,
-            0.8881359544145228,
-            0.45920099738999437,
-            -0.19284748723617093
-        ],
-        [
-            0.002854261818715398,
-            -0.4593266423892484,
-            0.8882628489252997,
-            0.057154511710361816
-        ],
-        [
-            0.9998214254141742,
-            -0.01528273756969839,
-            -0.011115539354645598,
-            -0.04187423675125192
-        ],
-        [
-            0.0,
-            0.0,
-            0.0,
-            1.0
-        ]
-    ]
-}

tools/add_n_txt.py

@@ -1,21 +0,0 @@
-import os
-
-# ======== 你可以修改的参数 ========
-start_id = 1           # 起始编号
-end_id = 200           # 结束编号
-save_dir = "labels"    # 输出目录
-prefix = "neg_"        # 文件名前缀
-zero_padding = 5       # 位数（例如 00001）
-# ==================================
-
-os.makedirs(save_dir, exist_ok=True)
-
-for i in range(start_id, end_id + 1):
-    name = f"{prefix}{str(i).zfill(zero_padding)}.txt"
-    path = os.path.join(save_dir, name)
-
-    # 创建空文件
-    with open(path, "w") as f:
-        pass
-
-    print("created:", path)

tools/calibration_tools.py

@@ -192,13 +192,13 @@ class _Calibration:
                     np.linalg.inv(
                         self._function(
                             self._hand[i], self._hand[i + 1], self._hand[i + 2],
-                            self._hand[i + 6], self._hand[i + 3], self._hand[i + 4], self._hand[i + 5]
+                            self._hand[i + 3], self._hand[i + 4], self._hand[i + 5], self._hand[i + 6]
                         )
                     )
                     if self._mode == 'eye_to_hand' else
                     self._function(
                         self._hand[i], self._hand[i + 1], self._hand[i + 2],
-                        self._hand[i + 6], self._hand[i + 3], self._hand[i + 4], self._hand[i + 5]
+                        self._hand[i + 3], self._hand[i + 4], self._hand[i + 5], self._hand[i + 6]
                     )
                 )
                 self._Hcs.append(

tools/cap_video.py

@@ -1,39 +0,0 @@
-import cv2
-# import numpy as np
-import time
-
-cap = cv2.VideoCapture(6)
-
-cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)   # 宽
-cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)   # 高
-
-frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
-fourcc = cv2.VideoWriter_fourcc(*'MP4V')  # 'XVID', 'MJPG', 'MP4V' 等
-out = cv2.VideoWriter('video/output_720p.mp4', fourcc, 20.0, (frame_width, frame_height))
-
-i = 0
-
-time.sleep(1)
-
-while True:
-    # 逐帧捕获
-    ret, frame = cap.read()
-
-    if not ret:
-        print("读取摄像头画面失败")
-        break
-
-    if i <= 5:
-        i = i+1
-        continue
-
-    cv2.imshow('camera', frame)
-    out.write(frame)
-    if cv2.waitKey(5) & 0xFF == ord('q'):
-        break
-
-cap.release()
-out.release()
-cv2.destroyAllWindows()

tools/crop_from_video.py

@@ -1,31 +0,0 @@
-import cv2
-import os
-
-video_path = "video/video_5.mp4"
-
-save_dir = "/home/lyx/Images/hivecore_box_datasets"
-os.makedirs(save_dir, exist_ok=True)
-
-cap = cv2.VideoCapture(video_path)
-if not cap.isOpened():
-    print("无法打开视频")
-    exit()
-
-frame_interval = 15
-frame_count = 0
-saved_count = 448
-
-while True:
-    ret, frame = cap.read()
-    if not ret:
-        break
-
-    if frame_count % frame_interval == 0:
-        save_path = os.path.join(save_dir, f"frame_{saved_count:06d}.jpg")
-        cv2.imwrite(save_path, frame)
-        saved_count += 1
-
-    frame_count += 1
-
-cap.release()
-print(f"共保存 {saved_count} 张图像到 {save_dir}")

tools/error_calculation_tool.py

@@ -52,7 +52,7 @@ def calculate(mode, hand_path, camera_path, hand_eye_path):
     for i in range(sum_):
         hand = [hand_all[i*7], hand_all[i*7+1], hand_all[i*7+2], hand_all[i*7+3], hand_all[i*7+4], hand_all[i*7+5], hand_all[i*7+6]]
         camera = [camera_all[i*7], camera_all[i*7+1], camera_all[i*7+2], camera_all[i*7+3], camera_all[i*7+4], camera_all[i*7+5], camera_all[i*7+6]]
-        T_hand_in_base = get_matrix_quat(hand[0], hand[1], hand[2], hand[6], hand[3], hand[4], hand[5])
+        T_hand_in_base = get_matrix_quat(hand[0], hand[1], hand[2], hand[3], hand[4], hand[5], hand[6])
         T_cal_in_camera = get_matrix_quat(camera[0], camera[1], camera[2], camera[3], camera[4], camera[5], camera[6])
 
         if mode == "eye_to_hand":

tools/get_file_name_list.py

@@ -1,16 +0,0 @@
-import os
-
-
-file_dir = "/home/lyx/Datasets/hivecore_box_datasets/"
-
-with open(os.path.join(file_dir, "train.txt"), "w") as f:
-    for file in os.listdir(os.path.join(file_dir, 'images/train/')):
-        if file.endswith(".jpg"):
-            labels_dir = os.path.join(os.path.join('./images/train/', file))
-            f.write(labels_dir + "\n")
-
-with open(os.path.join(file_dir, "val.txt"), "w") as f:
-    for file in os.listdir(os.path.join(file_dir, 'images/val/')):
-        if file.endswith(".jpg"):
-            labels_dir = os.path.join(os.path.join('./images/val/', file))
-            f.write(labels_dir + "\n")

tools/get_train_and_val.py

@@ -1,37 +0,0 @@
-import os
-import shutil
-
-
-def check_dir(dirname):
-    if not os.path.exists(dirname):
-        os.makedirs(dirname)
-
-
-file_dir = "/home/lyx/Datasets/hivecore_box_datasets/data"
-imgs_dir = "/home/lyx/Datasets/hivecore_box_datasets/images"
-labels_dir = "/home/lyx/Datasets/hivecore_box_datasets/labels"
-
-check_dir(imgs_dir)
-check_dir(os.path.join(imgs_dir, "train"))
-check_dir(os.path.join(imgs_dir, "val"))
-check_dir(labels_dir)
-check_dir(os.path.join(labels_dir, "train"))
-check_dir(os.path.join(labels_dir, "val"))
-
-i = 0
-for file in os.listdir(file_dir):
-    if file.endswith(".jpg"):
-        dot_index = file.rfind(".")
-        file_name = file[:dot_index]
-        label = file_name + '.txt'
-
-        if i % 10 == 9:
-            shutil.copy(os.path.join(file_dir, file), os.path.join(imgs_dir, "val",file))
-            shutil.copy(os.path.join(file_dir, label), os.path.join(labels_dir, "val", label))
-        else:
-            shutil.copy(os.path.join(file_dir, file), os.path.join(imgs_dir, "train", file))
-            shutil.copy(os.path.join(file_dir, label), os.path.join(labels_dir, "train", label))
-
-        i = i + 1
-
-

tools/json2yolo.py

@@ -1,56 +0,0 @@
-import json
-import os
- 
-label_to_class_id = {
-    "box": 0,  # 从0开始
-    # 其他类别...
-}
- 
- 
-def convert_labelme_json_to_yolo(json_file, output_dir):
-    try:
-        with open(json_file, 'r') as f:
-            labelme_data = json.load(f)
- 
-        img_width = labelme_data["imageWidth"]
-        img_height = labelme_data["imageHeight"]
- 
-        file_name = os.path.splitext(os.path.basename(json_file))[0]
-        txt_path = os.path.join(output_dir, f"{file_name}.txt")
- 
-        with open(txt_path, 'w') as txt_file:
-            for shape in labelme_data['shapes']:
-                label = shape['label']
-                points = shape['points']
- 
-                if not points:
-                    continue
- 
-                class_id = label_to_class_id.get(label)
-                if class_id is None:
-                    print(f"Warning: 跳过未定义标签 '{label}'")
-                    continue
- 
-                # 检查多边形是否闭合
-                if points[0] != points[-1]:
-                    points.append(points[0])
- 
-                normalized = [(x / img_width, y / img_height) for x, y in points]
-                line = f"{class_id} " + " ".join(f"{x:.6f} {y:.6f}" for x, y in normalized)
-                txt_file.write(line + "\n")
- 
-    except Exception as e:
-        print(f"处理文件 {json_file} 时出错: {str(e)}")
- 
- 
-if __name__ == "__main__":
-    json_dir = "/home/lyx/Datasets/hivecore_box_datasets/data"  # labelme标注存放的目录
-    output_dir = "/home/lyx/Datasets/hivecore_box_datasets/data"  # 输出目录
- 
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
- 
-    for json_file in os.listdir(json_dir):
-        if json_file.endswith(".json"):
-            json_path = os.path.join(json_dir, json_file)
-            convert_labelme_json_to_yolo(json_path, output_dir)

tools/pt2vino.py

@@ -1,24 +0,0 @@
-import argparse
-
-from ultralytics import YOLO
-
-
-def main(checkpoint_path):
-    model = YOLO(checkpoint_path)
-    model.export(
-        format="openvino",
-        imgsz=(1280, 720),
-        dynamic=True,
-        simplify=True,
-        half=True,
-        workspace=0.8,
-        batch=1,
-        device="cpu"
-    )
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("checkpoint_path", type=str)
-    args = parser.parse_args()
-    main(args.checkpoint_path)
-    # main(checkpoint_path="checkpoint/medical_sense-seg.pt")

vision_detect/configs/error_configs/report_logging_define.json

@@ -1,29 +0,0 @@
-{
-  "info": {},
-
-  "warring": {
-    "0000": "Success",
-
-    "1000": "Detected object count is 0",
-    "1001": "Depth crop is None",
-    "1003": "Failed to detect a valid pose",
-
-    "1100": "Object point cloud contains excessive noise",
-    "1101": "The point cloud is empty",
-
-    "1200": "The number of points is insufficient to compute an OBB",
-    "1201": "PCA output vector is None",
-    "1202": "This pose cannot be grab, and position refine fail",
-
-    "1300": "E2E model input data 'coors' are fewer than 128",
-    "1301": "E2E model input data 'point_clouds' are fewer than 128",
-    "1302": "The 'true num' of points is 0; No graspable points are available",
-    "1303": "The model returned no predictions",
-    "1304": "All rotation vector processing failed; no valid rotation matrix was generated"
-  },
-
-  "error": {},
-
-  "fatal": {}
-}
-

vision_detect/configs/launch_configs/bottle_detect_service_icp.json

@@ -2,31 +2,28 @@
     "node_name": "bottle_detect_service",
     "output_boxes": "True",
     "output_masks": "False",
-
-    "calibration": {
-        "left_hand": "configs/hand_eye_mat/eye_in_left_hand.json",
-        "right_hand": "configs/hand_eye_mat/eye_in_right_hand.json",
-        "head": "configs/hand_eye_mat/eye_to_hand.json"
-    },
-
+    "left_hand": "configs/hand_eye_mat/eye_in_left_hand.json",
+    "right_hand": "configs/hand_eye_mat/eye_in_right_hand.json",
+    "head": "configs/hand_eye_mat/eye_to_hand.json",
     "get_camera_mode": "Service",
     "Service_configs": {
         "service_name": "/vision_object_recognition"
     },
-
     "detect_mode": "Detect",
     "Detect_configs": {
         "checkpoint_path": "checkpoints/yolo11s-seg.pt",
         "confidence": 0.50,
         "classes": [39]
     },
-
     "calculate_mode": "PCA",
     "ICP_configs": {
         "complete_model_path": "pointclouds/bottle_model.pcd",
         "depth_scale": 1000.0,
         "depth_trunc": 2.0,
-        "voxel_size": 0.010,
+        "nb_points": 10,
+        "radius": 0.1,
+        "nb_neighbors": 20,
+        "std_ratio": 3.0,
         "ransac_voxel_size": 0.005,
         "icp_voxel_radius": [0.004, 0.002, 0.001],
         "icp_max_iter": [50, 30, 14]

vision_detect/configs/launch_configs/bottle_detect_service_pca.json

@@ -2,29 +2,27 @@
     "node_name": "bottle_detect_service",
     "output_boxes": "True",
     "output_masks": "False",
-
-    "calibration": {
-        "left_hand": "configs/hand_eye_mat/eye_in_left_hand.json",
-        "right_hand": "configs/hand_eye_mat/eye_in_right_hand.json",
-        "head": "configs/hand_eye_mat/eye_to_hand.json"
-    },
-
+    "left_hand": "configs/hand_eye_mat/eye_in_left_hand.json",
+    "right_hand": "configs/hand_eye_mat/eye_in_right_hand.json",
+    "head": "configs/hand_eye_mat/eye_to_hand.json",
     "get_camera_mode": "Service",
     "Service_configs": {
         "service_name": "/vision_object_recognition"
     },
-
     "detect_mode": "Detect",
     "Detect_configs": {
         "checkpoint_path": "checkpoints/yolo11s-seg.pt",
         "confidence": 0.50,
         "classes": [39]
     },
-
     "calculate_mode": "PCA",
     "PCA_configs": {
         "depth_scale": 1000.0,
         "depth_trunc": 3.0,
-        "voxel_size": 0.010
+        "voxel_size": 0.020,
+        "nb_points": 10,
+        "radius": 0.1,
+        "nb_neighbors": 20,
+        "std_ratio": 3.0
     }
 }

vision_detect/configs/launch_configs/crossboard_topic_pca.json

@@ -2,13 +2,9 @@
     "node_name": "crossboard_topic",
     "output_boxes": "True",
     "output_masks": "False",
-
-    "calibration": {
-        "left_hand": "configs/hand_eye_mat/eye_in_left_hand.json",
-        "right_hand": "configs/hand_eye_mat/eye_in_right_hand.json",
-        "head": "configs/hand_eye_mat/eye_to_hand.json"
-    },
-
+    "left_hand": "configs/hand_eye_mat/eye_in_left_hand.json",
+    "right_hand": "configs/hand_eye_mat/eye_in_right_hand.json",
+    "head": "configs/hand_eye_mat/eye_to_hand.json",
     "get_camera_mode": "Topic",
     "Topic_configs": {
         "position": "right",
@@ -16,12 +12,10 @@
         "depth_image_topic_name": "/camera/depth/image_raw",
         "camera_info_topic_name": "/camera/color/camera_info"
     },
-
     "detect_mode": "Crossboard",
     "Crossboard_configs": {
         "pattern_size": [8, 5]
     },
-
     "calculate_mode": "PCA",
     "PCA_configs": {
         "depth_scale": 1000.0,

vision_detect/configs/launch_configs/default_config.json

@@ -2,13 +2,9 @@
     "node_name": "default_config_detect_service",
     "output_boxes": "True",
     "output_masks": "False",
-
-    "calibration": {
-        "left_hand": "configs/hand_eye_mat/eye_in_left_hand.json",
-        "right_hand": "configs/hand_eye_mat/eye_in_right_hand.json",
-        "head": "configs/hand_eye_mat/eye_to_hand.json"
-    },
-
+    "left_hand": "configs/hand_eye_mat/eye_in_left_hand.json",
+    "right_hand": "configs/hand_eye_mat/eye_in_right_hand.json",
+    "head": "configs/hand_eye_mat/eye_to_hand.json",
     "get_camera_mode": "Service",
     "Service_configs": {
         "service_name": "/vision_object_recognition"
@@ -19,7 +15,6 @@
         "depth_image_topic_name": "/camera/depth/image_raw",
         "camera_info_topic_name": "/camera/color/camera_info"
     },
-
     "detect_mode": "Detect",
     "Detect_configs": {
         "checkpoint_path": "checkpoints/yolo11s-seg.pt",
@@ -33,27 +28,26 @@
     "Crossboard_configs": {
         "pattern_size": [8, 5]
     },
-
     "calculate_mode": "PCA",
     "PCA_configs": {
         "depth_scale": 1000.0,
         "depth_trunc": 3.0,
-        "voxel_size": 0.010
+        "voxel_size": 0.020,
+        "nb_points": 10,
+        "radius": 0.1,
+        "nb_neighbors": 20,
+        "std_ratio": 3.0
     },
     "ICP_configs": {
         "complete_model_path": "pointclouds/bottle_model.pcd",
         "depth_scale": 1000.0,
         "depth_trunc": 2.0,
-        "voxel_size": 0.010,
+        "nb_points": 10,
+        "radius": 0.1,
+        "nb_neighbors": 20,
+        "std_ratio": 3.0,
         "ransac_voxel_size": 0.005,
         "icp_voxel_radius": [0.004, 0.002, 0.001],
         "icp_max_iter": [50, 30, 14]
-    },
-    "E2E_configs": {
-        "checkpoint_path": "checkpoints/posenet.pt",
-        "depth_scale": 1000.0,
-        "depth_trunc": 3.0,
-        "voxel_size": 0.010,
-        "collision_thresh": 0.01
     }
 }

vision_detect/configs/launch_configs/detect_service_pca.json

@@ -2,29 +2,27 @@
     "node_name": "detect_service",
     "output_boxes": "True",
     "output_masks": "False",
-
-    "calibration": {
-        "left_hand": "configs/hand_eye_mat/eye_in_left_hand.json",
-        "right_hand": "configs/hand_eye_mat/eye_in_right_hand.json",
-        "head": "configs/hand_eye_mat/eye_to_hand.json"
-    },
-
+    "left_hand": "configs/hand_eye_mat/eye_in_left_hand.json",
+    "right_hand": "configs/hand_eye_mat/eye_in_right_hand.json",
+    "head": "configs/hand_eye_mat/eye_to_hand.json",
     "get_camera_mode": "Service",
     "Service_configs": {
         "service_name": "/vision_object_recognition"
     },
-
     "detect_mode": "Detect",
     "Detect_configs": {
         "checkpoint_path": "checkpoints/yolo11s-seg.pt",
         "confidence": 0.50,
         "classes": []
     },
-
     "calculate_mode": "PCA",
     "PCA_configs": {
         "depth_scale": 1000.0,
         "depth_trunc": 3.0,
-        "voxel_size": 0.010
+        "voxel_size": 0.020,
+        "nb_points": 10,
+        "radius": 0.1,
+        "nb_neighbors": 20,
+        "std_ratio": 3.0
     }
 }

vision_detect/configs/launch_configs/medical_sense.json

@@ -1,37 +0,0 @@
-{
-    "node_name": "medical_sense_service",
-    "output_boxes": "True",
-    "output_masks": "False",
-
-    "calibration": {
-        "left_hand": "configs/hand_eye_mat/eye_in_left_hand.json",
-        "right_hand": "configs/hand_eye_mat/eye_in_right_hand.json",
-        "head": "configs/hand_eye_mat/eye_to_hand.json"
-    },
-
-    "get_camera_mode": "Service",
-    "Service_configs": {
-        "service_name": "/vision_object_recognition"
-    },
-
-    "detect_mode": "Detect",
-    "Detect_configs": {
-        "checkpoint_path": "checkpoints/medical_sense-seg.pt",
-        "confidence": 0.75,
-        "classes": []
-    },
-
-    "calculate_mode": "PCA",
-    "PCA_configs": {
-        "depth_scale": 1000.0,
-        "depth_trunc": 3.0,
-        "voxel_size": 0.005
-    },
-    "E2E_configs": {
-        "checkpoint_path": "checkpoints/posenet.pt",
-        "depth_scale": 1000.0,
-        "depth_trunc": 3.0,
-        "voxel_size": 0.0005,
-        "collision_thresh": 0.01
-    }
-}

vision_detect/launch/bottle_detect_service_icp.launch.py

@@ -15,11 +15,35 @@ with open(config_dir, "r") as f:
 
 def generate_launch_description():
     args_detect = [
-        DeclareLaunchArgument('configs_path', default_value=config_dir),
+        DeclareLaunchArgument('output_boxes', default_value=configs['output_boxes']),
+        DeclareLaunchArgument('output_masks', default_value=configs['output_masks']),
+        DeclareLaunchArgument('left_hand', default_value=configs['left_hand']),
+        DeclareLaunchArgument('right_hand', default_value=configs['right_hand']),
+        DeclareLaunchArgument('head', default_value=configs['eye_to_handhead_mat_path']),
+
+        DeclareLaunchArgument('get_camera_mode', default_value=configs['get_camera_mode']),
+        DeclareLaunchArgument('detect_mode', default_value=configs['detect_mode']),
+        DeclareLaunchArgument('calculate_mode', default_value=configs['calculate_mode']),
+
+        DeclareLaunchArgument("Service_configs", default_value=json.dumps(configs['Service_configs'])),
+        DeclareLaunchArgument("Detect_configs", default_value=json.dumps(configs['Detect_configs'])),
+        DeclareLaunchArgument("ICP_configs", default_value=json.dumps(configs['ICP_configs'])),
     ]
 
     def create_detect_node(context):
-        configs_path = LaunchConfiguration('configs_path').perform(context)
+        output_boxes = LaunchConfiguration('output_boxes').perform(context)
+        output_masks = LaunchConfiguration('output_masks').perform(context)
+        left_hand = LaunchConfiguration('left_hand').perform(context)
+        right_hand = LaunchConfiguration('right_hand').perform(context)
+        head = LaunchConfiguration('head').perform(context)
+
+        get_camera_mode = LaunchConfiguration('get_camera_mode').perform(context)
+        detect_mode = LaunchConfiguration('detect_mode').perform(context)
+        calculate_mode = LaunchConfiguration('calculate_mode').perform(context)
+
+        Service_configs = LaunchConfiguration('Service_configs').perform(context)
+        Detect_configs = LaunchConfiguration('Detect_configs').perform(context)
+        ICP_configs = LaunchConfiguration('ICP_configs').perform(context)
 
         return [
             Node(
@@ -27,7 +51,19 @@ def generate_launch_description():
                 executable='detect_node',
                 name=configs['node_name'],
                 parameters=[{
-                    'configs_path': configs_path,
+                    'output_boxes': output_boxes.lower() == 'true',
+                    'output_masks': output_masks.lower() == 'true',
+                    'left_hand': left_hand,
+                    'right_hand': right_hand,
+                    'head': head,
+
+                    'get_camera_mode': get_camera_mode,
+                    'detect_mode': detect_mode,
+                    'calculate_mode': calculate_mode,
+
+                    'Service_configs': Service_configs,
+                    'Detect_configs': Detect_configs,
+                    'ICP_configs': ICP_configs,
                 }]
             )
         ]

vision_detect/launch/bottle_detect_service_pca.launch.py

@@ -15,11 +15,35 @@ with open(config_dir, "r") as f:
 
 def generate_launch_description():
     args_detect = [
-        DeclareLaunchArgument('configs_path', default_value=config_dir),
+        DeclareLaunchArgument('output_boxes', default_value=configs['output_boxes']),
+        DeclareLaunchArgument('output_masks', default_value=configs['output_masks']),
+        DeclareLaunchArgument('left_hand', default_value=configs['left_hand']),
+        DeclareLaunchArgument('right_hand', default_value=configs['right_hand']),
+        DeclareLaunchArgument('head', default_value=configs['head']),
+
+        DeclareLaunchArgument('get_camera_mode', default_value=configs['get_camera_mode']),
+        DeclareLaunchArgument('detect_mode', default_value=configs['detect_mode']),
+        DeclareLaunchArgument('calculate_mode', default_value=configs['calculate_mode']),
+
+        DeclareLaunchArgument("Service_configs", default_value=json.dumps(configs['Service_configs'])),
+        DeclareLaunchArgument("Detect_configs", default_value=json.dumps(configs['Detect_configs'])),
+        DeclareLaunchArgument("PCA_configs", default_value=json.dumps(configs['PCA_configs'])),
     ]
 
     def create_detect_node(context):
-        configs_path = LaunchConfiguration('configs_path').perform(context)
+        output_boxes = LaunchConfiguration('output_boxes').perform(context)
+        output_masks = LaunchConfiguration('output_masks').perform(context)
+        left_hand = LaunchConfiguration('left_hand').perform(context)
+        right_hand = LaunchConfiguration('right_hand').perform(context)
+        head = LaunchConfiguration('head').perform(context)
+
+        get_camera_mode = LaunchConfiguration('get_camera_mode').perform(context)
+        detect_mode = LaunchConfiguration('detect_mode').perform(context)
+        calculate_mode = LaunchConfiguration('calculate_mode').perform(context)
+
+        Service_configs = LaunchConfiguration('Service_configs').perform(context)
+        Detect_configs = LaunchConfiguration('Detect_configs').perform(context)
+        PCA_configs = LaunchConfiguration('PCA_configs').perform(context)
 
         return [
             Node(
@@ -27,7 +51,19 @@ def generate_launch_description():
                 executable='detect_node',
                 name=configs['node_name'],
                 parameters=[{
-                    'configs_path': configs_path,
+                    'output_boxes': output_boxes.lower() == 'true',
+                    'output_masks': output_masks.lower() == 'true',
+                    'left_hand': left_hand,
+                    'right_hand': right_hand,
+                    'head': head,
+
+                    'get_camera_mode': get_camera_mode,
+                    'detect_mode': detect_mode,
+                    'calculate_mode': calculate_mode,
+
+                    'Service_configs': Service_configs,
+                    'Detect_configs': Detect_configs,
+                    'PCA_configs': PCA_configs,
                 }]
             )
         ]

vision_detect/launch/crossboard.launch.py

@@ -0,0 +1,47 @@
+from launch import LaunchDescription
+from launch_ros.actions import Node
+from launch.actions import DeclareLaunchArgument, OpaqueFunction
+from launch.substitutions import LaunchConfiguration
+
+import os
+import ast
+import json
+from ament_index_python.packages import get_package_share_directory
+
+
+share_dir = get_package_share_directory('vision_detect')
+
+
+def generate_launch_description():
+    args_detect = [
+        DeclareLaunchArgument('output_boxes', default_value='True'),
+        DeclareLaunchArgument('output_masks', default_value='False'),
+        DeclareLaunchArgument('color_image_topic', default_value='/camera/color/image_raw'),
+        DeclareLaunchArgument('depth_image_topic', default_value='/camera/depth/image_raw'),
+        DeclareLaunchArgument('camera_info_topic', default_value='/camera/color/camera_info'),
+    ]
+
+    def create_detect_node(context):
+        output_boxes = LaunchConfiguration('output_boxes').perform(context)
+        output_masks = LaunchConfiguration('output_masks').perform(context)
+        color_image_topic = LaunchConfiguration('color_image_topic').perform(context)
+        depth_image_topic = LaunchConfiguration('depth_image_topic').perform(context)
+        camera_info_topic = LaunchConfiguration('camera_info_topic').perform(context)
+
+        return [
+            Node(
+                package='vision_detect',
+                executable='crossboard_detect_node',
+                parameters=[{
+                    'output_boxes': output_boxes.lower() == 'true',
+                    'output_masks': output_masks.lower() == 'true',
+                    'color_image_topic': color_image_topic,
+                    'depth_image_topic': depth_image_topic,
+                    'camera_info_topic': camera_info_topic,
+                }]
+            )
+        ]
+
+    return LaunchDescription(args_detect + [
+        OpaqueFunction(function=create_detect_node),
+    ])

vision_detect/launch/crossboard_topic_pca.launch.py

@@ -15,11 +15,36 @@ with open(config_dir, "r") as f:
 
 def generate_launch_description():
     args_detect = [
-        DeclareLaunchArgument('configs_path', default_value=config_dir),
+        DeclareLaunchArgument('output_boxes', default_value=configs['output_boxes']),
+        DeclareLaunchArgument('output_masks', default_value=configs['output_masks']),
+        DeclareLaunchArgument('left_hand', default_value=configs['left_hand']),
+        DeclareLaunchArgument('right_hand', default_value=configs['right_hand']),
+        DeclareLaunchArgument('head', default_value=configs['head']),
+
+
+        DeclareLaunchArgument('get_camera_mode', default_value=configs['get_camera_mode']),
+        DeclareLaunchArgument('detect_mode', default_value=configs['detect_mode']),
+        DeclareLaunchArgument('calculate_mode', default_value=configs['calculate_mode']),
+
+        DeclareLaunchArgument("Topic_configs", default_value=json.dumps(configs['Topic_configs'])),
+        DeclareLaunchArgument("Crossboard_configs", default_value=json.dumps(configs['Crossboard_configs'])),
+        DeclareLaunchArgument("PCA_configs", default_value=json.dumps(configs['PCA_configs'])),
     ]
 
     def create_detect_node(context):
-        configs_path = LaunchConfiguration('configs_path').perform(context)
+        output_boxes = LaunchConfiguration('output_boxes').perform(context)
+        output_masks = LaunchConfiguration('output_masks').perform(context)
+        left_hand = LaunchConfiguration('left_hand').perform(context)
+        right_hand = LaunchConfiguration('right_hand').perform(context)
+        head = LaunchConfiguration('head').perform(context)
+
+        get_camera_mode = LaunchConfiguration('get_camera_mode').perform(context)
+        detect_mode = LaunchConfiguration('detect_mode').perform(context)
+        calculate_mode = LaunchConfiguration('calculate_mode').perform(context)
+
+        Topic_configs = LaunchConfiguration('Topic_configs').perform(context)
+        Crossboard_configs = LaunchConfiguration('Crossboard_configs').perform(context)
+        PCA_configs = LaunchConfiguration('PCA_configs').perform(context)
 
         return [
             Node(
@@ -27,7 +52,19 @@ def generate_launch_description():
                 executable='detect_node',
                 name=configs['node_name'],
                 parameters=[{
-                    'configs_path': configs_path,
+                    'output_boxes': output_boxes.lower() == 'true',
+                    'output_masks': output_masks.lower() == 'true',
+                    'left_hand': left_hand,
+                    'right_hand': right_hand,
+                    'head': head,
+
+                    'get_camera_mode': get_camera_mode,
+                    'detect_mode': detect_mode,
+                    'calculate_mode': calculate_mode,
+
+                    'Topic_configs': Topic_configs,
+                    'Crossboard_configs': Crossboard_configs,
+                    'PCA_configs': PCA_configs,
                 }]
             )
         ]

vision_detect/launch/detect_service.launch.py

@@ -0,0 +1,49 @@
+from launch import LaunchDescription
+from launch_ros.actions import Node
+from launch.actions import DeclareLaunchArgument, OpaqueFunction
+from launch.substitutions import LaunchConfiguration
+
+import os
+import ast
+import json
+from ament_index_python.packages import get_package_share_directory
+
+
+share_dir = get_package_share_directory('vision_detect')
+config_dir = os.path.join(share_dir, 'configs/launch_configs/detect_service.json')
+with open(config_dir, "r") as f:
+    configs = json.load(f)
+
+def generate_launch_description():
+    args_detect = [
+        DeclareLaunchArgument('checkpoint_name', default_value=configs['checkpoint_name']),
+        DeclareLaunchArgument('output_boxes', default_value=configs['output_boxes']),
+        DeclareLaunchArgument('output_masks', default_value=configs['output_masks']),
+        DeclareLaunchArgument('set_confidence', default_value=configs['set_confidence']),
+        DeclareLaunchArgument('classes', default_value=configs['classes']),
+    ]
+
+    def create_detect_node(context):
+        checkpoint = LaunchConfiguration('checkpoint_name').perform(context)
+        output_boxes = LaunchConfiguration('output_boxes').perform(context)
+        output_masks = LaunchConfiguration('output_masks').perform(context)
+        conf = LaunchConfiguration('set_confidence').perform(context)
+        classes = LaunchConfiguration('classes').perform(context)
+
+        return [
+            Node(
+                package='vision_detect',
+                executable='detect_service_node',
+                parameters=[{
+                    'checkpoint_name': checkpoint,
+                    'output_boxes': output_boxes.lower() == 'true',
+                    'output_masks': output_masks.lower() == 'true',
+                    'set_confidence': float(conf),
+                    'classes': classes,
+                }]
+            )
+        ]
+
+    return LaunchDescription(args_detect + [
+        OpaqueFunction(function=create_detect_node),
+    ])

vision_detect/launch/detect_service_pca.launch.py

@@ -15,11 +15,35 @@ with open(config_dir, "r") as f:
 
 def generate_launch_description():
     args_detect = [
-        DeclareLaunchArgument('configs_path', default_value=config_dir),
+        DeclareLaunchArgument('output_boxes', default_value=configs['output_boxes']),
+        DeclareLaunchArgument('output_masks', default_value=configs['output_masks']),
+        DeclareLaunchArgument('left_hand', default_value=configs['left_hand']),
+        DeclareLaunchArgument('right_hand', default_value=configs['right_hand']),
+        DeclareLaunchArgument('head', default_value=configs['head']),
+
+        DeclareLaunchArgument('get_camera_mode', default_value=configs['get_camera_mode']),
+        DeclareLaunchArgument('detect_mode', default_value=configs['detect_mode']),
+        DeclareLaunchArgument('calculate_mode', default_value=configs['calculate_mode']),
+
+        DeclareLaunchArgument("Service_configs", default_value=json.dumps(configs['Service_configs'])),
+        DeclareLaunchArgument("Detect_configs", default_value=json.dumps(configs['Detect_configs'])),
+        DeclareLaunchArgument("PCA_configs", default_value=json.dumps(configs['PCA_configs'])),
     ]
 
     def create_detect_node(context):
-        configs_path = LaunchConfiguration('configs_path').perform(context)
+        output_boxes = LaunchConfiguration('output_boxes').perform(context)
+        output_masks = LaunchConfiguration('output_masks').perform(context)
+        left_hand = LaunchConfiguration('left_hand').perform(context)
+        right_hand = LaunchConfiguration('right_hand').perform(context)
+        head = LaunchConfiguration('head').perform(context)
+
+        get_camera_mode = LaunchConfiguration('get_camera_mode').perform(context)
+        detect_mode = LaunchConfiguration('detect_mode').perform(context)
+        calculate_mode = LaunchConfiguration('calculate_mode').perform(context)
+
+        Service_configs = LaunchConfiguration('Service_configs').perform(context)
+        Detect_configs = LaunchConfiguration('Detect_configs').perform(context)
+        PCA_configs = LaunchConfiguration('PCA_configs').perform(context)
 
         return [
             Node(
@@ -27,7 +51,19 @@ def generate_launch_description():
                 executable='detect_node',
                 name=configs['node_name'],
                 parameters=[{
-                    'configs_path': configs_path,
+                    'output_boxes': output_boxes.lower() == 'true',
+                    'output_masks': output_masks.lower() == 'true',
+                    'left_hand': left_hand,
+                    'right_hand': right_hand,
+                    'head': head,
+
+                    'get_camera_mode': get_camera_mode,
+                    'detect_mode': detect_mode,
+                    'calculate_mode': calculate_mode,
+
+                    'Service_configs': Service_configs,
+                    'Detect_configs': Detect_configs,
+                    'PCA_configs': PCA_configs,
                 }]
             )
         ]

vision_detect/launch/example_configs.launch.py

@@ -13,14 +13,74 @@ config_dir = os.path.join(share_dir, 'configs/launch_configs/default_config.json
 with open(config_dir, "r") as f:
     configs = json.load(f)
 
+# def get_param_configs():
+#     param_configs = {}
+#     if configs["get_camera_mode"] == "Service":
+#         param_configs.update(configs["Service_configs"])
+#     elif configs["get_camera_mode"] == "Topic":
+#         param_configs.update(configs["Topic_configs"])
+#     else:
+#         param_configs.update(configs["Service_configs"])
+#
+#     if configs["detect_mode"] == "Detect":
+#         param_configs.update(configs["Detect_configs"])
+#     elif configs["detect_mode"] == "Color":
+#         param_configs.update(configs["Color_configs"])
+#     elif configs["detect_mode"] == "Crossboard":
+#         param_configs.update(configs["Crossboard_configs"])
+#     else:
+#         param_configs.update(configs["Detect_configs"])
+#
+#     if configs["calculate_mode"] == "PCA":
+#         param_configs.update(configs["PCA_configs"])
+#     elif configs["calculate_mode"] == "ICP":
+#         param_configs.update(configs["ICP_configs"])
+#     else:
+#         param_configs.update(configs["PCA_configs"])
+#     return param_configs
+
 
 def generate_launch_description():
     args_detect = [
-        DeclareLaunchArgument('configs_path', default_value=config_dir),
+        DeclareLaunchArgument('output_boxes', default_value=configs['output_boxes']),
+        DeclareLaunchArgument('output_masks', default_value=configs['output_masks']),
+        DeclareLaunchArgument('left_hand', default_value=configs['left_hand']),
+        DeclareLaunchArgument('right_hand', default_value=configs['right_hand']),
+        DeclareLaunchArgument('head', default_value=configs['head']),
+
+        DeclareLaunchArgument('get_camera_mode', default_value=configs['get_camera_mode']),
+        DeclareLaunchArgument('detect_mode', default_value=configs['detect_mode']),
+        DeclareLaunchArgument('calculate_mode', default_value=configs['calculate_mode']),
+
+        # DeclareLaunchArgument("param_configs", default_value=json.dumps(get_param_configs())),
+        DeclareLaunchArgument("Service_configs", default_value=json.dumps(configs['Service_configs'])),
+        DeclareLaunchArgument("Topic_configs", default_value=json.dumps(configs['Topic_configs'])),
+        DeclareLaunchArgument("Detect_configs", default_value=json.dumps(configs['Detect_configs'])),
+        DeclareLaunchArgument("Color_configs", default_value=json.dumps(configs['Color_configs'])),
+        DeclareLaunchArgument("Crossboard_configs", default_value=json.dumps(configs['Crossboard_configs'])),
+        DeclareLaunchArgument("PCA_configs", default_value=json.dumps(configs['PCA_configs'])),
+        DeclareLaunchArgument("ICP_configs", default_value=json.dumps(configs['ICP_configs'])),
     ]
 
     def create_detect_node(context):
-        configs_path = LaunchConfiguration('configs_path').perform(context)
+        output_boxes = LaunchConfiguration('output_boxes').perform(context)
+        output_masks = LaunchConfiguration('output_masks').perform(context)
+        left_hand = LaunchConfiguration('left_hand').perform(context)
+        right_hand = LaunchConfiguration('right_hand').perform(context)
+        head = LaunchConfiguration('head').perform(context)
+
+        get_camera_mode = LaunchConfiguration('get_camera_mode').perform(context)
+        detect_mode = LaunchConfiguration('detect_mode').perform(context)
+        calculate_mode = LaunchConfiguration('calculate_mode').perform(context)
+
+        # param_configs = LaunchConfiguration('param_configs').perform(context)
+        Service_configs = LaunchConfiguration('Service_configs').perform(context)
+        Topic_configs = LaunchConfiguration('Topic_configs').perform(context)
+        Detect_configs = LaunchConfiguration('Detect_configs').perform(context)
+        Color_configs = LaunchConfiguration('Color_configs').perform(context)
+        Crossboard_configs = LaunchConfiguration('Crossboard_configs').perform(context)
+        PCA_configs = LaunchConfiguration('PCA_configs').perform(context)
+        ICP_configs = LaunchConfiguration('ICP_configs').perform(context)
 
         return [
             Node(
@@ -28,7 +88,24 @@ def generate_launch_description():
                 executable='detect_node',
                 name=configs['node_name'],
                 parameters=[{
-                    'configs_path': configs_path,
+                    'output_boxes': output_boxes.lower() == 'true',
+                    'output_masks': output_masks.lower() == 'true',
+                    'left_hand': left_hand,
+                    'right_hand': right_hand,
+                    'head': head,
+
+                    'get_camera_mode': get_camera_mode,
+                    'detect_mode': detect_mode,
+                    'calculate_mode': calculate_mode,
+
+                    # 'param_configs': param_configs,
+                    'Service_configs': Service_configs,
+                    'Topic_configs': Topic_configs,
+                    'Detect_configs': Detect_configs,
+                    'Color_configs': Color_configs,
+                    'Crossboard_configs': Crossboard_configs,
+                    'PCA_configs': PCA_configs,
+                    'ICP_configs': ICP_configs,
                 }]
             )
         ]

vision_detect/launch/medical_sense.launch.py

@@ -1,38 +0,0 @@
-from launch import LaunchDescription
-from launch_ros.actions import Node
-from launch.actions import DeclareLaunchArgument, OpaqueFunction
-from launch.substitutions import LaunchConfiguration
-
-import os
-import json
-from ament_index_python.packages import get_package_share_directory
-
-
-share_dir = get_package_share_directory('vision_detect')
-config_dir = os.path.join(share_dir, 'configs/launch_configs/medical_sense.json')
-with open(config_dir, "r") as f:
-    configs = json.load(f)
-
-def generate_launch_description():
-    args_detect = [
-        DeclareLaunchArgument('configs_path', default_value=config_dir),
-    ]
-
-    def create_detect_node(context):
-        configs_path = LaunchConfiguration('configs_path').perform(context)
-
-        return [
-            Node(
-                package='vision_detect',
-                executable='detect_node',
-                name=configs['node_name'],
-                output="screen",
-                parameters=[{
-                    'configs_path': configs_path,
-                }]
-            )
-        ]
-
-    return LaunchDescription(args_detect + [
-        OpaqueFunction(function=create_detect_node),
-    ])

vision_detect/setup.py

@@ -1,47 +1,26 @@
-import os
 from glob import glob
 from setuptools import find_packages, setup
 
 package_name = 'vision_detect'
 
-openvino_files = []
-for model_dir in glob("checkpoints/*_openvino_model"):
-    model_name = os.path.basename(model_dir)
-    model_files = glob(os.path.join(model_dir, '*'))
-    openvino_files.append(
-        ('share/' + package_name + '/checkpoints/' + model_name, model_files)
-    )
-
-data_files = [
-    ('share/ament_index/resource_index/packages', ['resource/' + package_name]),
-    ('share/' + package_name, ['package.xml']),
-    ('share/' + package_name + '/launch', glob('launch/*.launch.py')),
-    ('share/' + package_name + '/configs', glob('configs/*.json')),
-    ('share/' + package_name + '/configs/flexiv_configs',
-     glob('configs/flexiv_configs/*.json')),
-    ('share/' + package_name + '/configs/hand_eye_mat', glob('configs/hand_eye_mat/*.json')),
-    ('share/' + package_name + '/configs/launch_configs',
-     glob('configs/launch_configs/*.json')),
-    ('share/' + package_name + '/configs/error_configs', glob('configs/error_configs/*.json')),
-
-    ('share/' + package_name + '/checkpoints', glob('checkpoints/*.pt')),
-    ('share/' + package_name + '/checkpoints', glob('checkpoints/*.onnx')),
-    ('share/' + package_name + '/checkpoints', glob('checkpoints/*.engine')),
-    ('share/' + package_name + '/pointclouds', glob('pointclouds/*.pcd')),
-
-    ('lib/python3.10/site-packages/' + package_name + '/VisionDetect/net/pointnet2/pointnet2',
-     glob('vision_detect/VisionDetect/net/pointnet2/pointnet2/*.so')),
-    ('lib/python3.10/site-packages/' + package_name + '/VisionDetect/net/knn/knn_pytorch',
-     glob('vision_detect/VisionDetect/net/knn/knn_pytorch/*.so'))
-    ]
-
-data_files.extend(openvino_files)
-
 setup(
     name=package_name,
     version='0.0.0',
     packages=find_packages(exclude=['test']),
-    data_files=data_files,
+    data_files=[
+        ('share/ament_index/resource_index/packages', ['resource/' + package_name]),
+        ('share/' + package_name, ['package.xml']),
+        ('share/' + package_name + '/launch', glob('launch/*.launch.py')),
+        ('share/' + package_name + '/configs', glob('configs/*.json')),
+        ('share/' + package_name + '/configs/flexiv_configs', glob('configs/flexiv_configs/*.json')),
+        ('share/' + package_name + '/configs/hand_eye_mat', glob('configs/hand_eye_mat/*.json')),
+        ('share/' + package_name + '/configs/launch_configs', glob('configs/launch_configs/*.json')),
+
+        ('share/' + package_name + '/checkpoints', glob('checkpoints/*.pt')),
+        ('share/' + package_name + '/checkpoints', glob('checkpoints/*.onnx')),
+        ('share/' + package_name + '/checkpoints', glob('checkpoints/*.engine')),
+        ('share/' + package_name + '/pointclouds', glob('pointclouds/*.pcd')),
+    ],
     install_requires=['setuptools'],
     zip_safe=True,
     include_package_data=True,
@@ -52,15 +31,22 @@ setup(
     tests_require=['pytest'],
     entry_points={
         'console_scripts': [
-            # 'red_detect_topic_node = vision_detect.detect_red_topic:main',
-            # 'red_detect_service_node = vision_detect.detect_red_service:main',
+            'detect_service_node = vision_detect.detect_service:main',
+            'detect_topic_node = vision_detect.detect_topic:main',
+
+            'box_detect_service_node = vision_detect.detect_box_service:main',
+            'red_detect_topic_node = vision_detect.detect_red_topic:main',
+            'red_detect_service_node = vision_detect.detect_red_service:main',
 
             'flexiv_detect_topic_node = vision_detect.flexivaidk_detect_topic:main',
             'flexiv_detect_service_node = vision_detect.flexivaidk_detect_service:main',
 
             'sub_pose_node = vision_detect.sub_pose:main',
+            'calibration_node = vision_detect.hand_eye_calibration:main',
+            'crossboard_detect_node = vision_detect.crossboard_detect:main',
             'service_client_node = vision_detect.service_client:main',
             'get_camera_pose_node = vision_detect.get_camera_pose:main',
+            'calculate_node = vision_detect.calculate:main',
             
             'detect_node = vision_detect.detect_node:main',
         ],

vision_detect/vision_detect/VisionDetect/net/__init__.py

@@ -1 +0,0 @@
-from .gsnet import *

vision_detect/vision_detect/VisionDetect/net/gsnet/__init__.py

@@ -1,7 +0,0 @@
-from .models.graspnet import *
-from .utils.collision_detector import *
-
-
-__all__ = [
-    "GraspNet", "pred_decode", "ModelFreeCollisionDetector"
-]

vision_detect/vision_detect/VisionDetect/net/gsnet/gsnet_pose_calculate.py

@@ -1,132 +0,0 @@
-import os
-import sys
-import numpy as np
-import argparse
-from PIL import Image
-import torch
-import open3d as o3d
-
-ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(ROOT_DIR)
-
-from models.graspnet import GraspNet, pred_decode
-import collections.abc as container_abcs
-import MinkowskiEngine as ME
-from utils.collision_detector import ModelFreeCollisionDetector
-
-
-parser = argparse.ArgumentParser()
-parser.add_argument('--checkpoint_path', default='checkpoints/minkuresunet_realsense.tar')
-parser.add_argument('--dump_dir', help='Dump dir to save outputs', default='results/')
-parser.add_argument('--seed_feat_dim', default=512, type=int, help='Point wise feature dim')
-parser.add_argument('--collision_thresh', type=float, default=0.01,
-                    help='Collision Threshold in collision detection [default: 0.01]')
-parser.add_argument('--voxel_size_cd', type=float, default=0.01,
-                    help='Voxel Size for collision detection')
-parser.add_argument('--infer', action='store_true', default=True)
-parser.add_argument('--vis', action='store_true', default=True)
-cfgs = parser.parse_args()
-
-
-# ------------------------------------------------------------------------- GLOBAL CONFIG BEG
-# if not os.path.exists(cfgs.dump_dir):
-#     os.mkdir(cfgs.dump_dir)
-
-def minkowski_collate_fn(list_data):
-    coordinates_batch, features_batch = ME.utils.sparse_collate([d["coors"] for d in list_data],
-                                                                [d["feats"] for d in list_data])
-    coordinates_batch = np.ascontiguousarray(coordinates_batch, dtype=np.int32)
-    coordinates_batch, features_batch, _, quantize2original = ME.utils.sparse_quantize(
-        coordinates_batch, features_batch, return_index=True, return_inverse=True)
-    res = {
-        "coors": coordinates_batch,
-        "feats": features_batch,
-        "quantize2original": quantize2original
-    }
-
-    def collate_fn_(batch):
-        if type(batch[0]).__module__ == 'numpy':
-            return torch.stack([torch.from_numpy(b) for b in batch], 0)
-        elif isinstance(batch[0], container_abcs.Sequence):
-            return [[torch.from_numpy(sample) for sample in b] for b in batch]
-        elif isinstance(batch[0], container_abcs.Mapping):
-            for key in batch[0]:
-                if key == 'coors' or key == 'feats':
-                    continue
-                res[key] = collate_fn_([d[key] for d in batch])
-            return res
-    res = collate_fn_(list_data)
-
-    return res
-
-
-def data_process(pcd:o3d.geometry.PointCloud, voxel_size:float = 0.005):
-    index = 0
-    camera_poses = np.array([np.eye(4)]) # 相机位姿
-    align_mat = np.eye(4) # Camera_0 相对桌面的位姿
-    trans = np.dot(align_mat, camera_poses[int(index)])
-    pcd.transform(trans)
-    points = np.asarray(pcd.points)
-
-    ret_dict = {
-        'point_clouds': points.astype(np.float32),
-        'coors': points.astype(np.float32) / voxel_size,
-        'feats': np.ones_like(points).astype(np.float32),
-    }
-    return ret_dict
-
-
-def get_grasp_dict(preds):
-    grasp_dict = {
-            "score": preds[:, 0],
-            "width": preds[:, 1],
-            "height": preds[:, 2],
-            "depth": preds[:, 3],
-            "rotation": preds[:, 4:13].reshape(-1, 3, 3),
-            "translation": preds[:, 13:16].reshape(-1, 3),
-            "object_id": preds[:, 16]
-        }
-    return grasp_dict
-
-
-def inference(pcd:o3d.geometry.PointCloud, voxel_size:float = 0.005):
-    data_input = data_process(pcd, voxel_size)
-    batch_data = minkowski_collate_fn([data_input])
-    net = GraspNet(seed_feat_dim=512, is_training=False)
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    # device = torch.device("cpu")
-    net.to(device)
-    # Load checkpoint
-    checkpoint_path = os.path.join(ROOT_DIR, cfgs.checkpoint_path)
-    net.load_state_dict(checkpoint_path)
-
-    net.eval()
-
-    for key in batch_data:
-        batch_data[key] = batch_data[key].to(device)
-    # Forward pass
-    with torch.no_grad():
-        end_points = net(batch_data)
-        grasp_preds = pred_decode(end_points)
-
-    preds = grasp_preds[0].detach().cpu().numpy()
-
-    sorted_index = np.argsort(-preds[:, 0])
-    preds = preds[sorted_index]
-    preds = preds[:10]
-
-    # collision detection
-    if cfgs.collision_thresh > 0:
-        cloud = data_input['point_clouds']
-        mfcdetector = ModelFreeCollisionDetector(cloud, voxel_size=cfgs.voxel_size_cd)
-        collision_mask = mfcdetector.detect(get_grasp_dict(preds), approach_dist=0.05,
-                                            collision_thresh=cfgs.collision_thresh)
-        preds = preds[~collision_mask]
-
-    return preds
-
-
-if __name__ == '__main__':
-    data_dict = data_process()
-    inference(data_dict)
-

vision_detect/vision_detect/VisionDetect/net/gsnet/models/backbone_resunet14.py

@@ -1,224 +0,0 @@
-import MinkowskiEngine as ME
-from MinkowskiEngine.modules.resnet_block import BasicBlock, Bottleneck
-from .resnet import ResNetBase
-
-
-class MinkUNetBase(ResNetBase):
-    BLOCK = None
-    PLANES = None
-    DILATIONS = (1, 1, 1, 1, 1, 1, 1, 1)
-    LAYERS = (2, 2, 2, 2, 2, 2, 2, 2)
-    PLANES = (32, 64, 128, 256, 256, 128, 96, 96)
-    INIT_DIM = 32
-    OUT_TENSOR_STRIDE = 1
-
-    # To use the model, must call initialize_coords before forward pass.
-    # Once data is processed, call clear to reset the model before calling
-    # initialize_coords
-    def __init__(self, in_channels, out_channels, D=3):
-        ResNetBase.__init__(self, in_channels, out_channels, D)
-
-    def network_initialization(self, in_channels, out_channels, D):
-        # Output of the first conv concated to conv6
-        self.inplanes = self.INIT_DIM
-        self.conv0p1s1 = ME.MinkowskiConvolution(
-            in_channels, self.inplanes, kernel_size=5, dimension=D)
-
-        self.bn0 = ME.MinkowskiBatchNorm(self.inplanes)
-
-        self.conv1p1s2 = ME.MinkowskiConvolution(
-            self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D)
-        self.bn1 = ME.MinkowskiBatchNorm(self.inplanes)
-
-        self.block1 = self._make_layer(self.BLOCK, self.PLANES[0],
-                                       self.LAYERS[0])
-
-        self.conv2p2s2 = ME.MinkowskiConvolution(
-            self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D)
-        self.bn2 = ME.MinkowskiBatchNorm(self.inplanes)
-
-        self.block2 = self._make_layer(self.BLOCK, self.PLANES[1],
-                                       self.LAYERS[1])
-
-        self.conv3p4s2 = ME.MinkowskiConvolution(
-            self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D)
-
-        self.bn3 = ME.MinkowskiBatchNorm(self.inplanes)
-        self.block3 = self._make_layer(self.BLOCK, self.PLANES[2],
-                                       self.LAYERS[2])
-
-        self.conv4p8s2 = ME.MinkowskiConvolution(
-            self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D)
-        self.bn4 = ME.MinkowskiBatchNorm(self.inplanes)
-        self.block4 = self._make_layer(self.BLOCK, self.PLANES[3],
-                                       self.LAYERS[3])
-
-        self.convtr4p16s2 = ME.MinkowskiConvolutionTranspose(
-            self.inplanes, self.PLANES[4], kernel_size=2, stride=2, dimension=D)
-        self.bntr4 = ME.MinkowskiBatchNorm(self.PLANES[4])
-
-        self.inplanes = self.PLANES[4] + self.PLANES[2] * self.BLOCK.expansion
-        self.block5 = self._make_layer(self.BLOCK, self.PLANES[4],
-                                       self.LAYERS[4])
-        self.convtr5p8s2 = ME.MinkowskiConvolutionTranspose(
-            self.inplanes, self.PLANES[5], kernel_size=2, stride=2, dimension=D)
-        self.bntr5 = ME.MinkowskiBatchNorm(self.PLANES[5])
-
-        self.inplanes = self.PLANES[5] + self.PLANES[1] * self.BLOCK.expansion
-        self.block6 = self._make_layer(self.BLOCK, self.PLANES[5],
-                                       self.LAYERS[5])
-        self.convtr6p4s2 = ME.MinkowskiConvolutionTranspose(
-            self.inplanes, self.PLANES[6], kernel_size=2, stride=2, dimension=D)
-        self.bntr6 = ME.MinkowskiBatchNorm(self.PLANES[6])
-
-        self.inplanes = self.PLANES[6] + self.PLANES[0] * self.BLOCK.expansion
-        self.block7 = self._make_layer(self.BLOCK, self.PLANES[6],
-                                       self.LAYERS[6])
-        self.convtr7p2s2 = ME.MinkowskiConvolutionTranspose(
-            self.inplanes, self.PLANES[7], kernel_size=2, stride=2, dimension=D)
-        self.bntr7 = ME.MinkowskiBatchNorm(self.PLANES[7])
-
-        self.inplanes = self.PLANES[7] + self.INIT_DIM
-        self.block8 = self._make_layer(self.BLOCK, self.PLANES[7],
-                                       self.LAYERS[7])
-
-        self.final = ME.MinkowskiConvolution(
-            self.PLANES[7] * self.BLOCK.expansion,
-            out_channels,
-            kernel_size=1,
-            bias=True,
-            dimension=D)
-        self.relu = ME.MinkowskiReLU(inplace=True)
-
-    def forward(self, x):
-        out = self.conv0p1s1(x)
-        out = self.bn0(out)
-        out_p1 = self.relu(out)
-
-        out = self.conv1p1s2(out_p1)
-        out = self.bn1(out)
-        out = self.relu(out)
-        out_b1p2 = self.block1(out)
-
-        out = self.conv2p2s2(out_b1p2)
-        out = self.bn2(out)
-        out = self.relu(out)
-        out_b2p4 = self.block2(out)
-
-        out = self.conv3p4s2(out_b2p4)
-        out = self.bn3(out)
-        out = self.relu(out)
-        out_b3p8 = self.block3(out)
-
-        # tensor_stride=16
-        out = self.conv4p8s2(out_b3p8)
-        out = self.bn4(out)
-        out = self.relu(out)
-        out = self.block4(out)
-
-        # tensor_stride=8
-        out = self.convtr4p16s2(out)
-        out = self.bntr4(out)
-        out = self.relu(out)
-
-        out = ME.cat(out, out_b3p8)
-        out = self.block5(out)
-
-        # tensor_stride=4
-        out = self.convtr5p8s2(out)
-        out = self.bntr5(out)
-        out = self.relu(out)
-
-        out = ME.cat(out, out_b2p4)
-        out = self.block6(out)
-
-        # tensor_stride=2
-        out = self.convtr6p4s2(out)
-        out = self.bntr6(out)
-        out = self.relu(out)
-
-        out = ME.cat(out, out_b1p2)
-        out = self.block7(out)
-
-        # tensor_stride=1
-        out = self.convtr7p2s2(out)
-        out = self.bntr7(out)
-        out = self.relu(out)
-
-        out = ME.cat(out, out_p1)
-        out = self.block8(out)
-
-        return self.final(out)
-
-
-class MinkUNet14(MinkUNetBase):
-    BLOCK = BasicBlock
-    LAYERS = (1, 1, 1, 1, 1, 1, 1, 1)
-
-
-class MinkUNet18(MinkUNetBase):
-    BLOCK = BasicBlock
-    LAYERS = (2, 2, 2, 2, 2, 2, 2, 2)
-
-
-class MinkUNet34(MinkUNetBase):
-    BLOCK = BasicBlock
-    LAYERS = (2, 3, 4, 6, 2, 2, 2, 2)
-
-
-class MinkUNet50(MinkUNetBase):
-    BLOCK = Bottleneck
-    LAYERS = (2, 3, 4, 6, 2, 2, 2, 2)
-
-
-class MinkUNet101(MinkUNetBase):
-    BLOCK = Bottleneck
-    LAYERS = (2, 3, 4, 23, 2, 2, 2, 2)
-
-
-class MinkUNet14A(MinkUNet14):
-    PLANES = (32, 64, 128, 256, 128, 128, 96, 96)
-
-
-class MinkUNet14B(MinkUNet14):
-    PLANES = (32, 64, 128, 256, 128, 128, 128, 128)
-
-
-class MinkUNet14C(MinkUNet14):
-    PLANES = (32, 64, 128, 256, 192, 192, 128, 128)
-
-
-class MinkUNet14Dori(MinkUNet14):
-    PLANES = (32, 64, 128, 256, 384, 384, 384, 384)
-
-
-class MinkUNet14E(MinkUNet14):
-    PLANES = (32, 64, 128, 256, 384, 384, 384, 384)
-
-
-class MinkUNet14D(MinkUNet14):
-    PLANES = (32, 64, 128, 256, 192, 192, 192, 192)
-
-
-class MinkUNet18A(MinkUNet18):
-    PLANES = (32, 64, 128, 256, 128, 128, 96, 96)
-
-
-class MinkUNet18B(MinkUNet18):
-    PLANES = (32, 64, 128, 256, 128, 128, 128, 128)
-
-
-class MinkUNet18D(MinkUNet18):
-    PLANES = (32, 64, 128, 256, 384, 384, 384, 384)
-
-
-class MinkUNet34A(MinkUNet34):
-    PLANES = (32, 64, 128, 256, 256, 128, 64, 64)
-
-
-class MinkUNet34B(MinkUNet34):
-    PLANES = (32, 64, 128, 256, 256, 128, 64, 32)
-
-
-class MinkUNet34C(MinkUNet34):
-    PLANES = (32, 64, 128, 256, 256, 128, 96, 96)

vision_detect/vision_detect/VisionDetect/net/gsnet/models/graspnet.py

@@ -1,128 +0,0 @@
-""" GraspNet baseline model definition.
-    Author: chenxi-wang
-"""
-import logging
-
-# import os
-# import sys
-import numpy as np
-import torch
-import torch.nn as nn
-import MinkowskiEngine as ME
-
-from .backbone_resunet14 import MinkUNet14D
-from .modules import ApproachNet, GraspableNet, CloudCrop, SWADNet
-from ..utils.loss_utils import GRASP_MAX_WIDTH, NUM_VIEW, NUM_ANGLE, NUM_DEPTH, GRASPNESS_THRESHOLD, M_POINT
-from ..utils.label_generation import (process_grasp_labels, match_grasp_view_and_label, batch_viewpoint_params_to_matrix)
-from ...pointnet2.pointnet2_utils import furthest_point_sample, gather_operation
-
-
-class GraspNet(nn.Module):
-    def __init__(self, cylinder_radius=0.05, seed_feat_dim=512, is_training=True):
-        super().__init__()
-        self.is_training = is_training
-        self.seed_feature_dim = seed_feat_dim
-        self.num_depth = NUM_DEPTH
-        self.num_angle = NUM_ANGLE
-        self.M_points = M_POINT
-        self.num_view = NUM_VIEW
-
-        self.backbone = MinkUNet14D(in_channels=3, out_channels=self.seed_feature_dim, D=3)
-        self.graspable = GraspableNet(seed_feature_dim=self.seed_feature_dim)
-        self.rotation = ApproachNet(self.num_view, seed_feature_dim=self.seed_feature_dim, is_training=self.is_training)
-        self.crop = CloudCrop(nsample=16, cylinder_radius=cylinder_radius, seed_feature_dim=self.seed_feature_dim)
-        self.swad = SWADNet(num_angle=self.num_angle, num_depth=self.num_depth)
-
-    def forward(self, end_points):
-        seed_xyz = end_points['point_clouds']  # use all sampled point cloud, B*Ns*3
-        B, point_num, _ = seed_xyz.shape  # batch _size
-        # point-wise features
-        coordinates_batch = end_points['coors']
-        features_batch = end_points['feats']
-        mink_input = ME.SparseTensor(features_batch, coordinates=coordinates_batch)
-        seed_features = self.backbone(mink_input).F
-        seed_features = seed_features[end_points['quantize2original']].view(B, point_num, -1).transpose(1, 2)
-
-        end_points = self.graspable(seed_features, end_points)
-        seed_features_flipped = seed_features.transpose(1, 2)  # B*Ns*feat_dim
-        objectness_score = end_points['objectness_score']
-        graspness_score = end_points['graspness_score'].squeeze(1)
-        objectness_pred = torch.argmax(objectness_score, 1)
-        objectness_mask = (objectness_pred == 1)
-        graspness_mask = graspness_score > GRASPNESS_THRESHOLD
-        graspable_mask = objectness_mask & graspness_mask
-
-        seed_features_graspable = []
-        seed_xyz_graspable = []
-        graspable_num_batch = 0.
-        for i in range(B):
-            cur_mask = graspable_mask[i]
-            graspable_num_batch += cur_mask.sum()
-            cur_feat = seed_features_flipped[i][cur_mask]  # Ns*feat_dim
-            cur_seed_xyz = seed_xyz[i][cur_mask]  # Ns*3  Ns=0->error
-
-            logging.error(f"cur_mask: {cur_mask.shape}")
-            logging.error(f"true num: {torch.sum(cur_mask).item()}")
-            logging.error(f"cur_seed_xyz: {cur_seed_xyz.shape}")
-
-            if torch.sum(cur_mask).item() == 0:
-                return None
-
-            cur_seed_xyz = cur_seed_xyz.unsqueeze(0) # 1*Ns*3
-            fps_idxs = furthest_point_sample(cur_seed_xyz, self.M_points)
-            cur_seed_xyz_flipped = cur_seed_xyz.transpose(1, 2).contiguous()  # 1*3*Ns
-            cur_seed_xyz = gather_operation(cur_seed_xyz_flipped, fps_idxs).transpose(1, 2).squeeze(0).contiguous() # Ns*3
-            cur_feat_flipped = cur_feat.unsqueeze(0).transpose(1, 2).contiguous()  # 1*feat_dim*Ns
-            cur_feat = gather_operation(cur_feat_flipped, fps_idxs).squeeze(0).contiguous() # feat_dim*Ns
-
-            seed_features_graspable.append(cur_feat)
-            seed_xyz_graspable.append(cur_seed_xyz)
-
-        seed_xyz_graspable = torch.stack(seed_xyz_graspable, 0)  # B*Ns*3
-        seed_features_graspable = torch.stack(seed_features_graspable)  # B*feat_dim*Ns
-        end_points['xyz_graspable'] = seed_xyz_graspable
-        end_points['graspable_count_stage1'] = graspable_num_batch / B
-
-        end_points, res_feat = self.rotation(seed_features_graspable, end_points)
-        seed_features_graspable = seed_features_graspable + res_feat
-
-        if self.is_training:
-            end_points = process_grasp_labels(end_points)
-            grasp_top_views_rot, end_points = match_grasp_view_and_label(end_points)
-        else:
-            grasp_top_views_rot = end_points['grasp_top_view_rot']
-
-        group_features = self.crop(seed_xyz_graspable.contiguous(), seed_features_graspable.contiguous(), grasp_top_views_rot)
-        end_points = self.swad(group_features, end_points)
-
-        return end_points
-
-
-def pred_decode(end_points):
-    batch_size = len(end_points['point_clouds'])
-    grasp_preds = []
-    for i in range(batch_size):
-        grasp_center = end_points['xyz_graspable'][i].float()
-
-        grasp_score = end_points['grasp_score_pred'][i].float()
-        grasp_score = grasp_score.view(M_POINT, NUM_ANGLE*NUM_DEPTH)
-        grasp_score, grasp_score_inds = torch.max(grasp_score, -1)  # [M_POINT]
-        grasp_score = grasp_score.view(-1, 1)
-        grasp_angle = (grasp_score_inds // NUM_DEPTH) * np.pi / 12
-        grasp_depth = (grasp_score_inds % NUM_DEPTH + 1) * 0.01
-        grasp_depth = grasp_depth.view(-1, 1)
-        grasp_width = 1.2 * end_points['grasp_width_pred'][i] / 10.
-        grasp_width = grasp_width.view(M_POINT, NUM_ANGLE*NUM_DEPTH)
-        grasp_width = torch.gather(grasp_width, 1, grasp_score_inds.view(-1, 1))
-        grasp_width = torch.clamp(grasp_width, min=0., max=GRASP_MAX_WIDTH)
-
-        approaching = -end_points['grasp_top_view_xyz'][i].float()
-        grasp_rot = batch_viewpoint_params_to_matrix(approaching, grasp_angle)
-        grasp_rot = grasp_rot.view(M_POINT, 9)
-
-        # merge preds
-        grasp_height = 0.02 * torch.ones_like(grasp_score)
-        obj_ids = -1 * torch.ones_like(grasp_score)
-        grasp_preds.append(
-            torch.cat([grasp_score, grasp_width, grasp_height, grasp_depth, grasp_rot, grasp_center, obj_ids], axis=-1))
-    return grasp_preds

vision_detect/vision_detect/VisionDetect/net/gsnet/models/loss.py

@@ -1,80 +0,0 @@
-import torch.nn as nn
-import torch
-
-
-def get_loss(end_points):
-    objectness_loss, end_points = compute_objectness_loss(end_points)
-    graspness_loss, end_points = compute_graspness_loss(end_points)
-    view_loss, end_points = compute_view_graspness_loss(end_points)
-    score_loss, end_points = compute_score_loss(end_points)
-    width_loss, end_points = compute_width_loss(end_points)
-    loss = objectness_loss + 10 * graspness_loss + 100 * view_loss + 15 * score_loss + 10 * width_loss
-    end_points['loss/overall_loss'] = loss
-    return loss, end_points
-
-
-def compute_objectness_loss(end_points):
-    criterion = nn.CrossEntropyLoss(reduction='mean')
-    objectness_score = end_points['objectness_score']
-    objectness_label = end_points['objectness_label']
-    loss = criterion(objectness_score, objectness_label)
-    end_points['loss/stage1_objectness_loss'] = loss
-
-    objectness_pred = torch.argmax(objectness_score, 1)
-    end_points['stage1_objectness_acc'] = (objectness_pred == objectness_label.long()).float().mean()
-    end_points['stage1_objectness_prec'] = (objectness_pred == objectness_label.long())[
-        objectness_pred == 1].float().mean()
-    end_points['stage1_objectness_recall'] = (objectness_pred == objectness_label.long())[
-        objectness_label == 1].float().mean()
-    return loss, end_points
-
-
-def compute_graspness_loss(end_points):
-    criterion = nn.SmoothL1Loss(reduction='none')
-    graspness_score = end_points['graspness_score'].squeeze(1)
-    graspness_label = end_points['graspness_label'].squeeze(-1)
-    loss_mask = end_points['objectness_label'].bool()
-    loss = criterion(graspness_score, graspness_label)
-    loss = loss[loss_mask]
-    loss = loss.mean()
-    
-    graspness_score_c = graspness_score.detach().clone()[loss_mask]
-    graspness_label_c = graspness_label.detach().clone()[loss_mask]
-    graspness_score_c = torch.clamp(graspness_score_c, 0., 0.99)
-    graspness_label_c = torch.clamp(graspness_label_c, 0., 0.99)
-    rank_error = (torch.abs(torch.trunc(graspness_score_c * 20) - torch.trunc(graspness_label_c * 20)) / 20.).mean()
-    end_points['stage1_graspness_acc_rank_error'] = rank_error
-
-    end_points['loss/stage1_graspness_loss'] = loss
-    return loss, end_points
-
-
-def compute_view_graspness_loss(end_points):
-    criterion = nn.SmoothL1Loss(reduction='mean')
-    view_score = end_points['view_score']
-    view_label = end_points['batch_grasp_view_graspness']
-    loss = criterion(view_score, view_label)
-    end_points['loss/stage2_view_loss'] = loss
-    return loss, end_points
-
-
-def compute_score_loss(end_points):
-    criterion = nn.SmoothL1Loss(reduction='mean')
-    grasp_score_pred = end_points['grasp_score_pred']
-    grasp_score_label = end_points['batch_grasp_score']
-    loss = criterion(grasp_score_pred, grasp_score_label)
-
-    end_points['loss/stage3_score_loss'] = loss
-    return loss, end_points
-
-
-def compute_width_loss(end_points):
-    criterion = nn.SmoothL1Loss(reduction='none')
-    grasp_width_pred = end_points['grasp_width_pred']
-    grasp_width_label = end_points['batch_grasp_width'] * 10
-    loss = criterion(grasp_width_pred, grasp_width_label)
-    grasp_score_label = end_points['batch_grasp_score']
-    loss_mask = grasp_score_label > 0
-    loss = loss[loss_mask].mean()
-    end_points['loss/stage3_width_loss'] = loss
-    return loss, end_points

vision_detect/vision_detect/VisionDetect/net/gsnet/models/modules.py

@@ -1,116 +0,0 @@
-# import os
-# import sys
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-# BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-# ROOT_DIR = os.path.dirname(BASE_DIR)
-# sys.path.append(ROOT_DIR)
-
-from ...pointnet2 import pytorch_utils as pt_utils
-from ...pointnet2.pointnet2_utils import CylinderQueryAndGroup
-from ..utils.loss_utils import generate_grasp_views, batch_viewpoint_params_to_matrix
-
-
-class GraspableNet(nn.Module):
-    def __init__(self, seed_feature_dim):
-        super().__init__()
-        self.in_dim = seed_feature_dim  # default: 512
-        self.conv_graspable = nn.Conv1d(self.in_dim, 3, 1)
-
-    def forward(self, seed_features, end_points):
-        graspable_score = self.conv_graspable(seed_features)  # (B, 3, num_seed)
-        end_points['objectness_score'] = graspable_score[:, :2]  # (B, 3, num_seed)
-        end_points['graspness_score'] = graspable_score[:, 2]  # (B, 3, num_seed)
-        return end_points
-
-
-class ApproachNet(nn.Module):
-    def __init__(self, num_view, seed_feature_dim, is_training=True):
-        super().__init__()
-        self.num_view = num_view
-        self.in_dim = seed_feature_dim
-        self.is_training = is_training
-        self.conv1 = nn.Conv1d(self.in_dim, self.in_dim, 1)
-        self.conv2 = nn.Conv1d(self.in_dim, self.num_view, 1)
-
-    def forward(self, seed_features, end_points):
-        B, _, num_seed = seed_features.size()
-        res_features = F.relu(self.conv1(seed_features), inplace=True)
-        features = self.conv2(res_features)
-        view_score = features.transpose(1, 2).contiguous() # (B, num_seed, num_view)
-        end_points['view_score'] = view_score
-
-        if self.is_training:
-            # normalize view graspness score to 0~1
-            view_score_ = view_score.clone().detach()
-            view_score_max, _ = torch.max(view_score_, dim=2)
-            view_score_min, _ = torch.min(view_score_, dim=2)
-            view_score_max = view_score_max.unsqueeze(-1).expand(-1, -1, self.num_view)
-            view_score_min = view_score_min.unsqueeze(-1).expand(-1, -1, self.num_view)
-            view_score_ = (view_score_ - view_score_min) / (view_score_max - view_score_min + 1e-8)
-
-            top_view_inds = []
-            for i in range(B):
-                top_view_inds_batch = torch.multinomial(view_score_[i], 1, replacement=False)
-                top_view_inds.append(top_view_inds_batch)
-            top_view_inds = torch.stack(top_view_inds, dim=0).squeeze(-1)  # B, num_seed
-        else:
-            _, top_view_inds = torch.max(view_score, dim=2)  # (B, num_seed)
-
-            top_view_inds_ = top_view_inds.view(B, num_seed, 1, 1).expand(-1, -1, -1, 3).contiguous()
-            template_views = generate_grasp_views(self.num_view).to(features.device)  # (num_view, 3)
-            template_views = template_views.view(1, 1, self.num_view, 3).expand(B, num_seed, -1, -1).contiguous()
-            vp_xyz = torch.gather(template_views, 2, top_view_inds_).squeeze(2)  # (B, num_seed, 3)
-            vp_xyz_ = vp_xyz.view(-1, 3)
-            batch_angle = torch.zeros(vp_xyz_.size(0), dtype=vp_xyz.dtype, device=vp_xyz.device)
-            vp_rot = batch_viewpoint_params_to_matrix(-vp_xyz_, batch_angle).view(B, num_seed, 3, 3)
-            end_points['grasp_top_view_xyz'] = vp_xyz
-            end_points['grasp_top_view_rot'] = vp_rot
-
-        end_points['grasp_top_view_inds'] = top_view_inds
-        return end_points, res_features
-
-
-class CloudCrop(nn.Module):
-    def __init__(self, nsample, seed_feature_dim, cylinder_radius=0.05, hmin=-0.02, hmax=0.04):
-        super().__init__()
-        self.nsample = nsample
-        self.in_dim = seed_feature_dim
-        self.cylinder_radius = cylinder_radius
-        mlps = [3 + self.in_dim, 256, 256]   # use xyz, so plus 3
-
-        self.grouper = CylinderQueryAndGroup(radius=cylinder_radius, hmin=hmin, hmax=hmax, nsample=nsample,
-                                             use_xyz=True, normalize_xyz=True)
-        self.mlps = pt_utils.SharedMLP(mlps, bn=True)
-
-    def forward(self, seed_xyz_graspable, seed_features_graspable, vp_rot):
-        grouped_feature = self.grouper(seed_xyz_graspable, seed_xyz_graspable, vp_rot,
-                                       seed_features_graspable)  # B*3 + feat_dim*M*K
-        new_features = self.mlps(grouped_feature)  # (batch_size, mlps[-1], M, K)
-        new_features = F.max_pool2d(new_features, kernel_size=[1, new_features.size(3)])  # (batch_size, mlps[-1], M, 1)
-        new_features = new_features.squeeze(-1)   # (batch_size, mlps[-1], M)
-        return new_features
-
-
-class SWADNet(nn.Module):
-    def __init__(self, num_angle, num_depth):
-        super().__init__()
-        self.num_angle = num_angle
-        self.num_depth = num_depth
-
-        self.conv1 = nn.Conv1d(256, 256, 1)  # input feat dim need to be consistent with CloudCrop module
-        self.conv_swad = nn.Conv1d(256, 2*num_angle*num_depth, 1)
-
-    def forward(self, vp_features, end_points):
-        B, _, num_seed = vp_features.size()
-        vp_features = F.relu(self.conv1(vp_features), inplace=True)
-        vp_features = self.conv_swad(vp_features)
-        vp_features = vp_features.view(B, 2, self.num_angle, self.num_depth, num_seed)
-        vp_features = vp_features.permute(0, 1, 4, 2, 3)
-
-        # split prediction
-        end_points['grasp_score_pred'] = vp_features[:, 0]  # B * num_seed * num angle * num_depth
-        end_points['grasp_width_pred'] = vp_features[:, 1]
-        return end_points

vision_detect/vision_detect/VisionDetect/net/gsnet/models/resnet.py

@@ -1,196 +0,0 @@
-import torch.nn as nn
-
-try:
-    import open3d as o3d
-except ImportError:
-    raise ImportError("Please install open3d with `pip install open3d`.")
-
-import MinkowskiEngine as ME
-from MinkowskiEngine.modules.resnet_block import BasicBlock, Bottleneck
-
-
-class ResNetBase(nn.Module):
-    BLOCK = None
-    LAYERS = ()
-    INIT_DIM = 64
-    PLANES = (64, 128, 256, 512)
-
-    def __init__(self, in_channels, out_channels, D=3):
-        nn.Module.__init__(self)
-        self.D = D
-        assert self.BLOCK is not None
-
-        self.network_initialization(in_channels, out_channels, D)
-        self.weight_initialization()
-
-    def network_initialization(self, in_channels, out_channels, D):
-
-        self.inplanes = self.INIT_DIM
-        self.conv1 = nn.Sequential(
-            ME.MinkowskiConvolution(
-                in_channels, self.inplanes, kernel_size=3, stride=2, dimension=D
-            ),
-            ME.MinkowskiInstanceNorm(self.inplanes),
-            ME.MinkowskiReLU(inplace=True),
-            ME.MinkowskiMaxPooling(kernel_size=2, stride=2, dimension=D),
-        )
-
-        self.layer1 = self._make_layer(
-            self.BLOCK, self.PLANES[0], self.LAYERS[0], stride=2
-        )
-        self.layer2 = self._make_layer(
-            self.BLOCK, self.PLANES[1], self.LAYERS[1], stride=2
-        )
-        self.layer3 = self._make_layer(
-            self.BLOCK, self.PLANES[2], self.LAYERS[2], stride=2
-        )
-        self.layer4 = self._make_layer(
-            self.BLOCK, self.PLANES[3], self.LAYERS[3], stride=2
-        )
-
-        self.conv5 = nn.Sequential(
-            ME.MinkowskiDropout(),
-            ME.MinkowskiConvolution(
-                self.inplanes, self.inplanes, kernel_size=3, stride=3, dimension=D
-            ),
-            ME.MinkowskiInstanceNorm(self.inplanes),
-            ME.MinkowskiGELU(),
-        )
-
-        self.glob_pool = ME.MinkowskiGlobalMaxPooling()
-
-        self.final = ME.MinkowskiLinear(self.inplanes, out_channels, bias=True)
-
-    def weight_initialization(self):
-        for m in self.modules():
-            if isinstance(m, ME.MinkowskiConvolution):
-                ME.utils.kaiming_normal_(m.kernel, mode="fan_out", nonlinearity="relu")
-
-            if isinstance(m, ME.MinkowskiBatchNorm):
-                nn.init.constant_(m.bn.weight, 1)
-                nn.init.constant_(m.bn.bias, 0)
-
-    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, bn_momentum=0.1):
-        downsample = None
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            downsample = nn.Sequential(
-                ME.MinkowskiConvolution(
-                    self.inplanes,
-                    planes * block.expansion,
-                    kernel_size=1,
-                    stride=stride,
-                    dimension=self.D,
-                ),
-                ME.MinkowskiBatchNorm(planes * block.expansion),
-            )
-        layers = []
-        layers.append(
-            block(
-                self.inplanes,
-                planes,
-                stride=stride,
-                dilation=dilation,
-                downsample=downsample,
-                dimension=self.D,
-            )
-        )
-        self.inplanes = planes * block.expansion
-        for i in range(1, blocks):
-            layers.append(
-                block(
-                    self.inplanes, planes, stride=1, dilation=dilation, dimension=self.D
-                )
-            )
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x: ME.SparseTensor):
-        x = self.conv1(x)
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-        x = self.conv5(x)
-        x = self.glob_pool(x)
-        return self.final(x)
-
-
-class ResNet14(ResNetBase):
-    BLOCK = BasicBlock
-    LAYERS = (1, 1, 1, 1)
-
-
-class ResNet18(ResNetBase):
-    BLOCK = BasicBlock
-    LAYERS = (2, 2, 2, 2)
-
-
-class ResNet34(ResNetBase):
-    BLOCK = BasicBlock
-    LAYERS = (3, 4, 6, 3)
-
-
-class ResNet50(ResNetBase):
-    BLOCK = Bottleneck
-    LAYERS = (3, 4, 6, 3)
-
-
-class ResNet101(ResNetBase):
-    BLOCK = Bottleneck
-    LAYERS = (3, 4, 23, 3)
-
-
-class ResFieldNetBase(ResNetBase):
-    def network_initialization(self, in_channels, out_channels, D):
-        field_ch = 32
-        field_ch2 = 64
-        self.field_network = nn.Sequential(
-            ME.MinkowskiSinusoidal(in_channels, field_ch),
-            ME.MinkowskiBatchNorm(field_ch),
-            ME.MinkowskiReLU(inplace=True),
-            ME.MinkowskiLinear(field_ch, field_ch),
-            ME.MinkowskiBatchNorm(field_ch),
-            ME.MinkowskiReLU(inplace=True),
-            ME.MinkowskiToSparseTensor(),
-        )
-        self.field_network2 = nn.Sequential(
-            ME.MinkowskiSinusoidal(field_ch + in_channels, field_ch2),
-            ME.MinkowskiBatchNorm(field_ch2),
-            ME.MinkowskiReLU(inplace=True),
-            ME.MinkowskiLinear(field_ch2, field_ch2),
-            ME.MinkowskiBatchNorm(field_ch2),
-            ME.MinkowskiReLU(inplace=True),
-            ME.MinkowskiToSparseTensor(),
-        )
-
-        ResNetBase.network_initialization(self, field_ch2, out_channels, D)
-
-    def forward(self, x: ME.TensorField):
-        otensor = self.field_network(x)
-        otensor2 = self.field_network2(otensor.cat_slice(x))
-        return ResNetBase.forward(self, otensor2)
-
-
-class ResFieldNet14(ResFieldNetBase):
-    BLOCK = BasicBlock
-    LAYERS = (1, 1, 1, 1)
-
-
-class ResFieldNet18(ResFieldNetBase):
-    BLOCK = BasicBlock
-    LAYERS = (2, 2, 2, 2)
-
-
-class ResFieldNet34(ResFieldNetBase):
-    BLOCK = BasicBlock
-    LAYERS = (3, 4, 6, 3)
-
-
-class ResFieldNet50(ResFieldNetBase):
-    BLOCK = Bottleneck
-    LAYERS = (3, 4, 6, 3)
-
-
-class ResFieldNet101(ResFieldNetBase):
-    BLOCK = Bottleneck
-    LAYERS = (3, 4, 23, 3)

vision_detect/vision_detect/VisionDetect/net/gsnet/utils/collision_detector.py

@@ -1,132 +0,0 @@
-""" Collision detection to remove collided grasp pose predictions.
-Author: chenxi-wang
-"""
-
-# import os
-# import sys
-import numpy as np
-import open3d as o3d
-
-
-class ModelFreeCollisionDetector():
-    """ Collision detection in scenes without object labels. Current finger width and length are fixed.
-
-        Input:
-                scene_points: [numpy.ndarray, (N,3), numpy.float32]
-                    the scene points to detect
-                voxel_size: [float]
-                    used for downsample
-
-        Example usage:
-            mfcdetector = ModelFreeCollisionDetector(scene_points, voxel_size=0.005)
-            collision_mask = mfcdetector.detect(grasp_group, approach_dist=0.03)
-            collision_mask, iou_list = mfcdetector.detect(grasp_group, approach_dist=0.03, collision_thresh=0.05, return_ious=True)
-            collision_mask, empty_mask = mfcdetector.detect(grasp_group, approach_dist=0.03, collision_thresh=0.05,
-                                            return_empty_grasp=True, empty_thresh=0.01)
-            collision_mask, empty_mask, iou_list = mfcdetector.detect(grasp_group, approach_dist=0.03, collision_thresh=0.05,
-                                            return_empty_grasp=True, empty_thresh=0.01, return_ious=True)
-    """
-    def __init__(self, scene_points, voxel_size=0.005):
-        self.finger_width = 0.005
-        self.finger_length = 0.09
-        self.voxel_size = voxel_size
-        scene_cloud = o3d.geometry.PointCloud()
-        scene_cloud.points = o3d.utility.Vector3dVector(scene_points)
-        scene_cloud = scene_cloud.voxel_down_sample(voxel_size)
-        self.scene_points = np.array(scene_cloud.points) # (K, 3)
-
-    def detect(self, preds, approach_dist=0.03, collision_thresh=0.05, return_empty_grasp=False, empty_thresh=0.01, return_ious=False):
-        """ Detect collision of grasps.
-
-            Input:
-                grasp_group: [GraspGroup, M grasps]
-                    the grasps to check
-                approach_dist: [float]
-                    the distance for a gripper to move along approaching direction before grasping
-                    this shifting space requires no voxel_sizepoint either
-                collision_thresh: [float]
-                    if global collision iou is greater than this threshold,
-                    a collision is detected
-                return_empty_grasp: [bool]
-                    if True, return a mask to imply whether there are objects in a grasp
-                empty_thresh: [float]
-                    if inner space iou is smaller than this threshold,
-                    a collision is detected
-                    only set when [return_empty_grasp] is True
-                return_ious: [bool]
-                    if True, return global collision iou and part collision ious
-                    
-            Output:
-                collision_mask: [numpy.ndarray, (M,), numpy.bool]
-                    True implies collision
-                [optional] empty_mask: [numpy.ndarray, (M,), numpy.bool]
-                    True implies empty grasp
-                    only returned when [return_empty_grasp] is True
-                [optional] iou_list: list of [numpy.ndarray, (M,), numpy.float32]
-                    global and part collision ious, containing
-                    [global_iou, left_iou, right_iou, bottom_iou, shifting_iou]
-                    only returned when [return_ious] is True
-        """
-        approach_dist = max(approach_dist, self.finger_width)
-        T = preds[:, 13:16].reshape(-1, 3) # (N, 3)
-        R = preds[:, 4:13].reshape(-1, 3, 3) # (N, 3, 3)
-        heights = preds[:, 2][:,np.newaxis] # (N,) -> (N, 1)
-        depths = preds[:, 3][:,np.newaxis] # (N,) -> (N, 1)
-        widths = preds[:, 1][:,np.newaxis] # (N,) -> (N, 1)
-        # 使用不同的抓取点偏置
-        targets = self.scene_points[np.newaxis,:,:] - T[:,np.newaxis,:] # (1, K, 3) - (N, 1, 3) = (N, K, 3)
-
-        targets = np.matmul(targets, R) # (N, K, 3) * (N, 3, 3) = (N, K, 3)
-
-        ## collision detection
-        # height mask
-        mask1 = ((targets[:,:,2] > -heights/2) & (targets[:,:,2] < heights/2)) # (N, K) * (N, 3) = (N, K)
-        # left finger mask
-        mask2 = ((targets[:,:,0] > depths - self.finger_length) & (targets[:,:,0] < depths))
-        mask3 = (targets[:,:,1] > -(widths/2 + self.finger_width))
-        mask4 = (targets[:,:,1] < -widths/2)
-        # right finger mask
-        mask5 = (targets[:,:,1] < (widths/2 + self.finger_width))
-        mask6 = (targets[:,:,1] > widths/2)
-        # bottom mask
-        mask7 = ((targets[:,:,0] <= depths - self.finger_length)\
-                & (targets[:,:,0] > depths - self.finger_length - self.finger_width))
-        # shifting mask
-        mask8 = ((targets[:,:,0] <= depths - self.finger_length - self.finger_width)\
-                & (targets[:,:,0] > depths - self.finger_length - self.finger_width - approach_dist))
-
-        # get collision mask of each point
-        left_mask = (mask1 & mask2 & mask3 & mask4)
-        right_mask = (mask1 & mask2 & mask5 & mask6)
-        bottom_mask = (mask1 & mask3 & mask5 & mask7)
-        shifting_mask = (mask1 & mask3 & mask5 & mask8)
-        global_mask = (left_mask | right_mask | bottom_mask | shifting_mask)
-
-        # calculate equivalant volume of each part
-        left_right_volume = (heights * self.finger_length * self.finger_width / (self.voxel_size**3)).reshape(-1)
-        bottom_volume = (heights * (widths+2*self.finger_width) * self.finger_width / (self.voxel_size**3)).reshape(-1)
-        shifting_volume = (heights * (widths+2*self.finger_width) * approach_dist / (self.voxel_size**3)).reshape(-1)
-        volume = left_right_volume*2 + bottom_volume + shifting_volume
-
-        # get collision iou of each part
-        global_iou = global_mask.sum(axis=1) / (volume+1e-6)
-
-        # get collison mask
-        collision_mask = (global_iou > collision_thresh)
-
-        if not (return_empty_grasp or return_ious):
-            return collision_mask
-
-        ret_value = [collision_mask,]
-        if return_empty_grasp:
-            inner_mask = (mask1 & mask2 & (~mask4) & (~mask6))
-            inner_volume = (heights * self.finger_length * widths / (self.voxel_size**3)).reshape(-1)
-            empty_mask = (inner_mask.sum(axis=-1)/inner_volume < empty_thresh)
-            ret_value.append(empty_mask)
-        if return_ious:
-            left_iou = left_mask.sum(axis=1) / (left_right_volume+1e-6)
-            right_iou = right_mask.sum(axis=1) / (left_right_volume+1e-6)
-            bottom_iou = bottom_mask.sum(axis=1) / (bottom_volume+1e-6)
-            shifting_iou = shifting_mask.sum(axis=1) / (shifting_volume+1e-6)
-            ret_value.append([global_iou, left_iou, right_iou, bottom_iou, shifting_iou])
-        return ret_value

vision_detect/vision_detect/VisionDetect/net/gsnet/utils/data_utils.py

@@ -1,156 +0,0 @@
-""" Tools for data processing.
-    Author: chenxi-wang
-"""
-
-import numpy as np
-
-
-class CameraInfo():
-    """ Camera intrisics for point cloud creation. """
-
-    def __init__(self, width, height, fx, fy, cx, cy, scale):
-        self.width = width
-        self.height = height
-        self.fx = fx
-        self.fy = fy
-        self.cx = cx
-        self.cy = cy
-        self.scale = scale
-
-
-def create_point_cloud_from_depth_image(depth, camera, organized=True):
-    """ Generate point cloud using depth image only.
-
-        Input:
-            depth: [numpy.ndarray, (H,W), numpy.float32]
-                depth image
-            camera: [CameraInfo]
-                camera intrinsics
-            organized: bool
-                whether to keep the cloud in image shape (H,W,3)
-
-        Output:
-            cloud: [numpy.ndarray, (H,W,3)/(H*W,3), numpy.float32]
-                generated cloud, (H,W,3) for organized=True, (H*W,3) for organized=False
-    """
-    assert (depth.shape[0] == camera.height and depth.shape[1] == camera.width)
-    xmap = np.arange(camera.width)
-    ymap = np.arange(camera.height)
-    xmap, ymap = np.meshgrid(xmap, ymap)
-    points_z = depth / camera.scale
-    points_x = (xmap - camera.cx) * points_z / camera.fx
-    points_y = (ymap - camera.cy) * points_z / camera.fy
-    cloud = np.stack([points_x, points_y, points_z], axis=-1)
-    if not organized:
-        cloud = cloud.reshape([-1, 3])
-    return cloud
-
-
-def transform_point_cloud(cloud, transform, format='4x4'):
-    """ Transform points to new coordinates with transformation matrix.
-
-        Input:
-            cloud: [np.ndarray, (N,3), np.float32]
-                points in original coordinates
-            transform: [np.ndarray, (3,3)/(3,4)/(4,4), np.float32]
-                transformation matrix, could be rotation only or rotation+translation
-            format: [string, '3x3'/'3x4'/'4x4']
-                the shape of transformation matrix
-                '3x3' --> rotation matrix
-                '3x4'/'4x4' --> rotation matrix + translation matrix
-
-        Output:
-            cloud_transformed: [np.ndarray, (N,3), np.float32]
-                points in new coordinates
-    """
-    if not (format == '3x3' or format == '4x4' or format == '3x4'):
-        raise ValueError('Unknown transformation format, only support \'3x3\' or \'4x4\' or \'3x4\'.')
-    if format == '3x3':
-        cloud_transformed = np.dot(transform, cloud.T).T
-    elif format == '4x4' or format == '3x4':
-        ones = np.ones(cloud.shape[0])[:, np.newaxis]
-        cloud_ = np.concatenate([cloud, ones], axis=1)
-        cloud_transformed = np.dot(transform, cloud_.T).T
-        cloud_transformed = cloud_transformed[:, :3]
-    return cloud_transformed
-
-
-def compute_point_dists(A, B):
-    """ Compute pair-wise point distances in two matrices.
-
-        Input:
-            A: [np.ndarray, (N,3), np.float32]
-                point cloud A
-            B: [np.ndarray, (M,3), np.float32]
-                point cloud B
-
-        Output:
-            dists: [np.ndarray, (N,M), np.float32]
-                distance matrix
-    """
-    A = A[:, np.newaxis, :]
-    B = B[np.newaxis, :, :]
-    dists = np.linalg.norm(A - B, axis=-1)
-    return dists
-
-
-def remove_invisible_grasp_points(cloud, grasp_points, pose, th=0.01):
-    """ Remove invisible part of object model according to scene point cloud.
-
-        Input:
-            cloud: [np.ndarray, (N,3), np.float32]
-                scene point cloud
-            grasp_points: [np.ndarray, (M,3), np.float32]
-                grasp point label in object coordinates
-            pose: [np.ndarray, (4,4), np.float32]
-                transformation matrix from object coordinates to world coordinates
-            th: [float]
-                if the minimum distance between a grasp point and the scene points is greater than outlier, the point will be removed
-
-        Output:
-            visible_mask: [np.ndarray, (M,), np.bool]
-                mask to show the visible part of grasp points
-    """
-    grasp_points_trans = transform_point_cloud(grasp_points, pose)
-    dists = compute_point_dists(grasp_points_trans, cloud)
-    min_dists = dists.min(axis=1)
-    visible_mask = (min_dists < th)
-    return visible_mask
-
-
-def get_workspace_mask(cloud, seg, trans=None, organized=True, outlier=0):
-    """ Keep points in workspace as input.
-
-        Input:
-            cloud: [np.ndarray, (H,W,3), np.float32]
-                scene point cloud
-            seg: [np.ndarray, (H,W,), np.uint8]
-                segmantation label of scene points
-            trans: [np.ndarray, (4,4), np.float32]
-                transformation matrix for scene points, default: None.
-            organized: [bool]
-                whether to keep the cloud in image shape (H,W,3)
-            outlier: [float]
-                if the distance between a point and workspace is greater than outlier, the point will be removed
-                
-        Output:
-            workspace_mask: [np.ndarray, (H,W)/(H*W,), np.bool]
-                mask to indicate whether scene points are in workspace
-    """
-    if organized:
-        h, w, _ = cloud.shape
-        cloud = cloud.reshape([h * w, 3])
-        seg = seg.reshape(h * w)
-    if trans is not None:
-        cloud = transform_point_cloud(cloud, trans)
-    foreground = cloud[seg > 0]
-    xmin, ymin, zmin = foreground.min(axis=0)
-    xmax, ymax, zmax = foreground.max(axis=0)
-    mask_x = ((cloud[:, 0] > xmin - outlier) & (cloud[:, 0] < xmax + outlier))
-    mask_y = ((cloud[:, 1] > ymin - outlier) & (cloud[:, 1] < ymax + outlier))
-    mask_z = ((cloud[:, 2] > zmin - outlier) & (cloud[:, 2] < zmax + outlier))
-    workspace_mask = (mask_x & mask_y & mask_z)
-    if organized:
-        workspace_mask = workspace_mask.reshape([h, w])
-
-    return workspace_mask

vision_detect/vision_detect/VisionDetect/net/gsnet/utils/label_generation.py

@@ -1,143 +0,0 @@
-""" Dynamically generate grasp labels during training.
-    Author: chenxi-wang
-"""
-
-import os
-import sys
-import torch
-
-# BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-# ROOT_DIR = os.path.dirname(BASE_DIR)
-# sys.path.append(ROOT_DIR)
-# sys.path.append(os.path.join(ROOT_DIR, 'knn'))
-
-from ...knn.knn_modules import knn
-from ..utils.loss_utils import GRASP_MAX_WIDTH, batch_viewpoint_params_to_matrix, \
-    transform_point_cloud, generate_grasp_views
-
-
-def process_grasp_labels(end_points):
-    """ Process labels according to scene points and object poses. """
-    seed_xyzs = end_points['xyz_graspable']  # (B, M_point, 3)
-    batch_size, num_samples, _ = seed_xyzs.size()
-
-    batch_grasp_points = []
-    batch_grasp_views_rot = []
-    batch_grasp_scores = []
-    batch_grasp_widths = []
-    for i in range(batch_size):
-        seed_xyz = seed_xyzs[i]  # (Ns, 3)
-        poses = end_points['object_poses_list'][i]  # [(3, 4),]
-
-        # get merged grasp points for label computation
-        grasp_points_merged = []
-        grasp_views_rot_merged = []
-        grasp_scores_merged = []
-        grasp_widths_merged = []
-        for obj_idx, pose in enumerate(poses):
-            grasp_points = end_points['grasp_points_list'][i][obj_idx]  # (Np, 3)
-            grasp_scores = end_points['grasp_scores_list'][i][obj_idx]  # (Np, V, A, D)
-            grasp_widths = end_points['grasp_widths_list'][i][obj_idx]  # (Np, V, A, D)
-            _, V, A, D = grasp_scores.size()
-            num_grasp_points = grasp_points.size(0)
-            # generate and transform template grasp views
-            grasp_views = generate_grasp_views(V).to(pose.device)  # (V, 3)
-            grasp_points_trans = transform_point_cloud(grasp_points, pose, '3x4')
-            grasp_views_trans = transform_point_cloud(grasp_views, pose[:3, :3], '3x3')
-            # generate and transform template grasp view rotation
-            angles = torch.zeros(grasp_views.size(0), dtype=grasp_views.dtype, device=grasp_views.device)
-            grasp_views_rot = batch_viewpoint_params_to_matrix(-grasp_views, angles)  # (V, 3, 3)
-            grasp_views_rot_trans = torch.matmul(pose[:3, :3], grasp_views_rot)  # (V, 3, 3)
-
-            # assign views
-            grasp_views_ = grasp_views.transpose(0, 1).contiguous().unsqueeze(0)
-            grasp_views_trans_ = grasp_views_trans.transpose(0, 1).contiguous().unsqueeze(0)
-            view_inds = knn(grasp_views_trans_, grasp_views_, k=1).squeeze() - 1
-            grasp_views_rot_trans = torch.index_select(grasp_views_rot_trans, 0, view_inds)  # (V, 3, 3)
-            grasp_views_rot_trans = grasp_views_rot_trans.unsqueeze(0).expand(num_grasp_points, -1, -1,
-                                                                              -1)  # (Np, V, 3, 3)
-            grasp_scores = torch.index_select(grasp_scores, 1, view_inds)  # (Np, V, A, D)
-            grasp_widths = torch.index_select(grasp_widths, 1, view_inds)  # (Np, V, A, D)
-            # add to list
-            grasp_points_merged.append(grasp_points_trans)
-            grasp_views_rot_merged.append(grasp_views_rot_trans)
-            grasp_scores_merged.append(grasp_scores)
-            grasp_widths_merged.append(grasp_widths)
-
-        grasp_points_merged = torch.cat(grasp_points_merged, dim=0)  # (Np', 3)
-        grasp_views_rot_merged = torch.cat(grasp_views_rot_merged, dim=0)  # (Np', V, 3, 3)
-        grasp_scores_merged = torch.cat(grasp_scores_merged, dim=0)  # (Np', V, A, D)
-        grasp_widths_merged = torch.cat(grasp_widths_merged, dim=0)  # (Np', V, A, D)
-
-        # compute nearest neighbors
-        seed_xyz_ = seed_xyz.transpose(0, 1).contiguous().unsqueeze(0)  # (1, 3, Ns)
-        grasp_points_merged_ = grasp_points_merged.transpose(0, 1).contiguous().unsqueeze(0)  # (1, 3, Np')
-        nn_inds = knn(grasp_points_merged_, seed_xyz_, k=1).squeeze() - 1  # (Ns)
-
-        # assign anchor points to real points
-        grasp_points_merged = torch.index_select(grasp_points_merged, 0, nn_inds)  # (Ns, 3)
-        grasp_views_rot_merged = torch.index_select(grasp_views_rot_merged, 0, nn_inds)  # (Ns, V, 3, 3)
-        grasp_scores_merged = torch.index_select(grasp_scores_merged, 0, nn_inds)  # (Ns, V, A, D)
-        grasp_widths_merged = torch.index_select(grasp_widths_merged, 0, nn_inds)  # (Ns, V, A, D)
-
-        # add to batch
-        batch_grasp_points.append(grasp_points_merged)
-        batch_grasp_views_rot.append(grasp_views_rot_merged)
-        batch_grasp_scores.append(grasp_scores_merged)
-        batch_grasp_widths.append(grasp_widths_merged)
-
-    batch_grasp_points = torch.stack(batch_grasp_points, 0)  # (B, Ns, 3)
-    batch_grasp_views_rot = torch.stack(batch_grasp_views_rot, 0)  # (B, Ns, V, 3, 3)
-    batch_grasp_scores = torch.stack(batch_grasp_scores, 0)  # (B, Ns, V, A, D)
-    batch_grasp_widths = torch.stack(batch_grasp_widths, 0)  # (B, Ns, V, A, D)
-
-    # compute view graspness
-    view_u_threshold = 0.6
-    view_grasp_num = 48
-    batch_grasp_view_valid_mask = (batch_grasp_scores <= view_u_threshold) & (batch_grasp_scores > 0) # (B, Ns, V, A, D)
-    batch_grasp_view_valid = batch_grasp_view_valid_mask.float()
-    batch_grasp_view_graspness = torch.sum(torch.sum(batch_grasp_view_valid, dim=-1), dim=-1) / view_grasp_num  # (B, Ns, V)
-    view_graspness_min, _ = torch.min(batch_grasp_view_graspness, dim=-1)  # (B, Ns)
-    view_graspness_max, _ = torch.max(batch_grasp_view_graspness, dim=-1)
-    view_graspness_max = view_graspness_max.unsqueeze(-1).expand(-1, -1, 300)  # (B, Ns, V)
-    view_graspness_min = view_graspness_min.unsqueeze(-1).expand(-1, -1, 300)  # same shape as batch_grasp_view_graspness
-    batch_grasp_view_graspness = (batch_grasp_view_graspness - view_graspness_min) / (view_graspness_max - view_graspness_min + 1e-5)
-
-    # process scores
-    label_mask = (batch_grasp_scores > 0) & (batch_grasp_widths <= GRASP_MAX_WIDTH)  # (B, Ns, V, A, D)
-    batch_grasp_scores[~label_mask] = 0
-
-    end_points['batch_grasp_point'] = batch_grasp_points
-    end_points['batch_grasp_view_rot'] = batch_grasp_views_rot
-    end_points['batch_grasp_score'] = batch_grasp_scores
-    end_points['batch_grasp_width'] = batch_grasp_widths
-    end_points['batch_grasp_view_graspness'] = batch_grasp_view_graspness
-
-    return end_points
-
-
-def match_grasp_view_and_label(end_points):
-    """ Slice grasp labels according to predicted views. """
-    top_view_inds = end_points['grasp_top_view_inds']  # (B, Ns)
-    template_views_rot = end_points['batch_grasp_view_rot']  # (B, Ns, V, 3, 3)
-    grasp_scores = end_points['batch_grasp_score']  # (B, Ns, V, A, D)
-    grasp_widths = end_points['batch_grasp_width']  # (B, Ns, V, A, D, 3)
-
-    B, Ns, V, A, D = grasp_scores.size()
-    top_view_inds_ = top_view_inds.view(B, Ns, 1, 1, 1).expand(-1, -1, -1, 3, 3)
-    top_template_views_rot = torch.gather(template_views_rot, 2, top_view_inds_).squeeze(2)
-    top_view_inds_ = top_view_inds.view(B, Ns, 1, 1, 1).expand(-1, -1, -1, A, D)
-    top_view_grasp_scores = torch.gather(grasp_scores, 2, top_view_inds_).squeeze(2)
-    top_view_grasp_widths = torch.gather(grasp_widths, 2, top_view_inds_).squeeze(2)
-
-    u_max = top_view_grasp_scores.max()
-    po_mask = top_view_grasp_scores > 0
-    po_mask_num = torch.sum(po_mask)
-    if po_mask_num > 0:
-        u_min = top_view_grasp_scores[po_mask].min()
-        top_view_grasp_scores[po_mask] = torch.log(u_max / top_view_grasp_scores[po_mask]) / (torch.log(u_max / u_min) + 1e-6)
-
-    end_points['batch_grasp_score'] = top_view_grasp_scores  # (B, Ns, A, D)
-    end_points['batch_grasp_width'] = top_view_grasp_widths  # (B, Ns, A, D)
-
-    return top_template_views_rot, end_points

vision_detect/vision_detect/VisionDetect/net/gsnet/utils/loss_utils.py

@@ -1,121 +0,0 @@
-""" Tools for loss computation.
-    Author: chenxi-wang
-"""
-
-import torch
-import numpy as np
-
-GRASP_MAX_WIDTH = 0.1
-GRASPNESS_THRESHOLD = 0.1
-NUM_VIEW = 300
-NUM_ANGLE = 12
-NUM_DEPTH = 4
-M_POINT = 1024
-
-
-def transform_point_cloud(cloud, transform, format='4x4'):
-    """ Transform points to new coordinates with transformation matrix.
-
-        Input:
-            cloud: [torch.FloatTensor, (N,3)]
-                points in original coordinates
-            transform: [torch.FloatTensor, (3,3)/(3,4)/(4,4)]
-                transformation matrix, could be rotation only or rotation+translation
-            format: [string, '3x3'/'3x4'/'4x4']
-                the shape of transformation matrix
-                '3x3' --> rotation matrix
-                '3x4'/'4x4' --> rotation matrix + translation matrix
-
-        Output:
-            cloud_transformed: [torch.FloatTensor, (N,3)]
-                points in new coordinates
-    """
-    if not (format == '3x3' or format == '4x4' or format == '3x4'):
-        raise ValueError('Unknown transformation format, only support \'3x3\' or \'4x4\' or \'3x4\'.')
-    if format == '3x3':
-        cloud_transformed = torch.matmul(transform, cloud.T).T
-    elif format == '4x4' or format == '3x4':
-        ones = cloud.new_ones(cloud.size(0), device=cloud.device).unsqueeze(-1)
-        cloud_ = torch.cat([cloud, ones], dim=1)
-        cloud_transformed = torch.matmul(transform, cloud_.T).T
-        cloud_transformed = cloud_transformed[:, :3]
-    return cloud_transformed
-
-
-def generate_grasp_views(N=300, phi=(np.sqrt(5) - 1) / 2, center=np.zeros(3), r=1):
-    """ View sampling on a unit sphere using Fibonacci lattices.
-        Ref: https://arxiv.org/abs/0912.4540
-
-        Input:
-            N: [int]
-                number of sampled views
-            phi: [float]
-                constant for view coordinate calculation, different phi's bring different distributions, default: (sqrt(5)-1)/2
-            center: [np.ndarray, (3,), np.float32]
-                sphere center
-            r: [float]
-                sphere radius
-
-        Output:
-            views: [torch.FloatTensor, (N,3)]
-                sampled view coordinates
-    """
-    views = []
-    for i in range(N):
-        zi = (2 * i + 1) / N - 1
-        xi = np.sqrt(1 - zi ** 2) * np.cos(2 * i * np.pi * phi)
-        yi = np.sqrt(1 - zi ** 2) * np.sin(2 * i * np.pi * phi)
-        views.append([xi, yi, zi])
-    views = r * np.array(views) + center
-    return torch.from_numpy(views.astype(np.float32))
-
-
-def batch_viewpoint_params_to_matrix(batch_towards, batch_angle):
-    """ Transform approach vectors and in-plane rotation angles to rotation matrices.
-
-        Input:
-            batch_towards: [torch.FloatTensor, (N,3)]
-                approach vectors in batch
-            batch_angle: [torch.floatTensor, (N,)]
-                in-plane rotation angles in batch
-                
-        Output:
-            batch_matrix: [torch.floatTensor, (N,3,3)]
-                rotation matrices in batch
-    """
-    axis_x = batch_towards
-    ones = torch.ones(axis_x.shape[0], dtype=axis_x.dtype, device=axis_x.device)
-    zeros = torch.zeros(axis_x.shape[0], dtype=axis_x.dtype, device=axis_x.device)
-    axis_y = torch.stack([-axis_x[:, 1], axis_x[:, 0], zeros], dim=-1)
-    mask_y = (torch.norm(axis_y, dim=-1) == 0)
-    axis_y[mask_y, 1] = 1
-    axis_x = axis_x / torch.norm(axis_x, dim=-1, keepdim=True)
-    axis_y = axis_y / torch.norm(axis_y, dim=-1, keepdim=True)
-    axis_z = torch.cross(axis_x, axis_y)
-    sin = torch.sin(batch_angle)
-    cos = torch.cos(batch_angle)
-    R1 = torch.stack([ones, zeros, zeros, zeros, cos, -sin, zeros, sin, cos], dim=-1)
-    R1 = R1.reshape([-1, 3, 3])
-    R2 = torch.stack([axis_x, axis_y, axis_z], dim=-1)
-    batch_matrix = torch.matmul(R2, R1)
-    return batch_matrix
-
-
-def huber_loss(error, delta=1.0):
-    """
-    Args:
-        error: Torch tensor (d1,d2,...,dk)
-    Returns:
-        loss: Torch tensor (d1,d2,...,dk)
-
-    x = error = pred - gt or dist(pred,gt)
-    0.5 * |x|^2                 if |x|<=d
-    0.5 * d^2 + d * (|x|-d)     if |x|>d
-    Author: Charles R. Qi
-    Ref: https://github.com/charlesq34/frustum-pointnets/blob/master/models/model_util.py
-    """
-    abs_error = torch.abs(error)
-    quadratic = torch.clamp(abs_error, max=delta)
-    linear = (abs_error - quadratic)
-    loss = 0.5 * quadratic ** 2 + delta * linear
-    return loss

vision_detect/vision_detect/VisionDetect/net/knn/knn_modules.py

@@ -1,20 +0,0 @@
-# import unittest
-# import gc
-# import operator as op
-# import functools
-import torch
-# from torch.autograd import Variable, Function
-# from knn_pytorch import knn_pytorch
-# import knn_pytorch
-from . import knn_pytorch
-
-
-def knn(ref, query, k=1):
-  """ Compute k nearest neighbors for each query point.
-  """
-  device = ref.device
-  ref = ref.float().to(device)
-  query = query.float().to(device)
-  inds = torch.empty(query.shape[0], k, query.shape[2]).long().to(device)
-  knn_pytorch.knn(ref, query, inds)
-  return inds

vision_detect/vision_detect/VisionDetect/net/knn/knn_pytorch/__init__.py

@@ -1 +0,0 @@
-from .knn_pytorch import *

vision_detect/vision_detect/VisionDetect/net/knn/setup.py

@@ -1,66 +0,0 @@
-#!/usr/bin/env python
-
-import glob
-import os
-
-import torch
-from setuptools import find_packages
-from setuptools import setup
-from torch.utils.cpp_extension import CUDA_HOME
-from torch.utils.cpp_extension import CppExtension
-from torch.utils.cpp_extension import CUDAExtension
-
-requirements = ["torch", "torchvision"]
-
-
-def get_extensions():
-    this_dir = os.path.dirname(os.path.abspath(__file__))
-    extensions_dir = os.path.join(this_dir, "src")
-
-    main_file = glob.glob(os.path.join(extensions_dir, "*.cpp"))
-    source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp"))
-    source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu"))
-
-    sources = main_file + source_cpu
-    extension = CppExtension
-
-    extra_compile_args = {"cxx": []}
-    define_macros = []
-
-    if torch.cuda.is_available() and CUDA_HOME is not None:
-        extension = CUDAExtension
-        sources += source_cuda
-        define_macros += [("WITH_CUDA", None)]
-        extra_compile_args["nvcc"] = [
-            "-DCUDA_HAS_FP16=1",
-            "-D__CUDA_NO_HALF_OPERATORS__",
-            "-D__CUDA_NO_HALF_CONVERSIONS__",
-            "-D__CUDA_NO_HALF2_OPERATORS__",
-        ]
-
-    sources = [os.path.join(extensions_dir, s) for s in sources]
-
-    include_dirs = [extensions_dir]
-
-    ext_modules = [
-        extension(
-            "knn_pytorch.knn_pytorch",
-            sources,
-            include_dirs=include_dirs,
-            define_macros=define_macros,
-            extra_compile_args=extra_compile_args,
-        )
-    ]
-
-    return ext_modules
-
-
-setup(
-    name="knn_pytorch",
-    version="0.1",
-    author="foolyc",
-    url="https://github.com/foolyc/torchKNN",
-    description="KNN implement in Pytorch 1.0 including both cpu version and gpu version",
-    ext_modules=get_extensions(),
-    cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension},
-)

vision_detect/vision_detect/VisionDetect/net/knn/src/cpu/knn_cpu.cpp

@@ -1,56 +0,0 @@
-#include "cpu/vision.h"
-
-
-void knn_cpu(float* ref_dev, int ref_width, float* query_dev, int query_width,
-    int height, int k, float* dist_dev, long* ind_dev, long* ind_buf)
-{
-    // Compute all the distances
-    for(int query_idx = 0;query_idx<query_width;query_idx++)
-    {
-        for(int ref_idx = 0;ref_idx < ref_width;ref_idx++)
-        {
-            dist_dev[query_idx * ref_width + ref_idx] = 0;
-            for(int hi=0;hi<height;hi++)
-                dist_dev[query_idx * ref_width + ref_idx] += (ref_dev[hi * ref_width + ref_idx] - query_dev[hi * query_width + query_idx]) * (ref_dev[hi * ref_width + ref_idx] - query_dev[hi * query_width + query_idx]);
-        }
-    }
-
-    float temp_value;
-    long temp_idx;
-    // sort the distance and get the index
-    for(int query_idx = 0;query_idx<query_width;query_idx++)
-    {
-        for(int i = 0;i < ref_width;i++)
-        {
-            ind_buf[i] = i+1;
-        }
-        for(int i = 0;i < ref_width;i++)
-            for(int j = 0;j < ref_width - i - 1;j++)
-            {
-                if(dist_dev[query_idx * ref_width + j] > dist_dev[query_idx * ref_width + j + 1])
-                {
-                    temp_value = dist_dev[query_idx * ref_width + j];
-                    dist_dev[query_idx * ref_width + j] = dist_dev[query_idx * ref_width + j + 1];
-                    dist_dev[query_idx * ref_width + j + 1] = temp_value;
-                    temp_idx = ind_buf[j];
-                    ind_buf[j] = ind_buf[j + 1];
-                    ind_buf[j + 1] = temp_idx;
-                }
-
-            }
-
-        for(int i = 0;i < k;i++)
-            ind_dev[query_idx + i * query_width] = ind_buf[i];
-        #if DEBUG
-        for(int i = 0;i < ref_width;i++)
-            printf("%d, ", ind_buf[i]);
-        printf("\n");
-        #endif
-
-    }
-
-
-
-
-
-}

vision_detect/vision_detect/VisionDetect/net/knn/src/cpu/vision.h

@@ -1,6 +0,0 @@
-#pragma once
-#include <torch/extension.h>
-
-void knn_cpu(float* ref_dev, int ref_width,
-    float* query_dev, int query_width,
-    int height, int k, float* dist_dev, long* ind_dev, long* ind_buf);

vision_detect/vision_detect/VisionDetect/net/knn/src/cuda/knn.cu

@@ -1,269 +0,0 @@
-/** Modifed version of knn-CUDA from https://github.com/vincentfpgarcia/kNN-CUDA
- * The modifications are
- *      removed texture memory usage
- *      removed split query KNN computation
- *      added feature extraction with bilinear interpolation
- *
- * Last modified by Christopher B. Choy <chrischoy@ai.stanford.edu> 12/23/2016
- */
-
-// Includes
-#include <cstdio>
-#include "cuda.h"
-
-#define IDX2D(i, j, dj) (dj * i + j)
-#define IDX3D(i, j, k, dj, dk) (IDX2D(IDX2D(i, j, dj), k, dk))
-
-#define BLOCK 512
-#define MAX_STREAMS 512
-
-// Constants used by the program
-#define BLOCK_DIM                      16
-#define DEBUG                          0
-
-
-/**
-  * Computes the distance between two matrix A (reference points) and
-  * B (query points) containing respectively wA and wB points.
-  *
-  * @param A     pointer on the matrix A
-  * @param wA    width of the matrix A = number of points in A
-  * @param B     pointer on the matrix B
-  * @param wB    width of the matrix B = number of points in B
-  * @param dim   dimension of points = height of matrices A and B
-  * @param AB    pointer on the matrix containing the wA*wB distances computed
-  */
-__global__ void cuComputeDistanceGlobal( float* A, int wA,
-    float* B, int wB, int dim, float* AB){
-
-// Declaration of the shared memory arrays As and Bs used to store the sub-matrix of A and B
-__shared__ float shared_A[BLOCK_DIM][BLOCK_DIM];
-__shared__ float shared_B[BLOCK_DIM][BLOCK_DIM];
-
-
-  // Sub-matrix of A (begin, step, end) and Sub-matrix of B (begin, step)
-  __shared__ int begin_A;
-  __shared__ int begin_B;
-  __shared__ int step_A;
-  __shared__ int step_B;
-  __shared__ int end_A;
-
-  // Thread index
-  int tx = threadIdx.x;
-  int ty = threadIdx.y;
-
-  // Other variables
-  float tmp;
-  float ssd = 0;
-
-  // Loop parameters
-  begin_A = BLOCK_DIM * blockIdx.y;
-  begin_B = BLOCK_DIM * blockIdx.x;
-  step_A  = BLOCK_DIM * wA;
-  step_B  = BLOCK_DIM * wB;
-  end_A   = begin_A + (dim-1) * wA;
-
-    // Conditions
-  int cond0 = (begin_A + tx < wA); // used to write in shared memory
-  int cond1 = (begin_B + tx < wB); // used to write in shared memory & to computations and to write in output matrix
-  int cond2 = (begin_A + ty < wA); // used to computations and to write in output matrix
-
-  // Loop over all the sub-matrices of A and B required to compute the block sub-matrix
-  for (int a = begin_A, b = begin_B; a <= end_A; a += step_A, b += step_B) {
-    // Load the matrices from device memory to shared memory; each thread loads one element of each matrix
-    if (a/wA + ty < dim){
-      shared_A[ty][tx] = (cond0)? A[a + wA * ty + tx] : 0;
-      shared_B[ty][tx] = (cond1)? B[b + wB * ty + tx] : 0;
-    }
-    else{
-      shared_A[ty][tx] = 0;
-      shared_B[ty][tx] = 0;
-    }
-
-    // Synchronize to make sure the matrices are loaded
-    __syncthreads();
-
-    // Compute the difference between the two matrixes; each thread computes one element of the block sub-matrix
-    if (cond2 && cond1){
-      for (int k = 0; k < BLOCK_DIM; ++k){
-        tmp = shared_A[k][ty] - shared_B[k][tx];
-        ssd += tmp*tmp;
-      }
-    }
-
-    // Synchronize to make sure that the preceding computation is done before loading two new sub-matrices of A and B in the next iteration
-    __syncthreads();
-  }
-
-  // Write the block sub-matrix to device memory; each thread writes one element
-  if (cond2 && cond1)
-    AB[(begin_A + ty) * wB + begin_B + tx] = ssd;
-}
-
-
-/**
-  * Gathers k-th smallest distances for each column of the distance matrix in the top.
-  *
-  * @param dist        distance matrix
-  * @param ind         index matrix
-  * @param width       width of the distance matrix and of the index matrix
-  * @param height      height of the distance matrix and of the index matrix
-  * @param k           number of neighbors to consider
-  */
-__global__ void cuInsertionSort(float *dist, long *ind, int width, int height, int k){
-
-  // Variables
-  int l, i, j;
-  float *p_dist;
-  long  *p_ind;
-  float curr_dist, max_dist;
-  long  curr_row,  max_row;
-  unsigned int xIndex = blockIdx.x * blockDim.x + threadIdx.x;
-
-  if (xIndex<width){
-    // Pointer shift, initialization, and max value
-    p_dist   = dist + xIndex;
-    p_ind    = ind  + xIndex;
-    max_dist = p_dist[0];
-    p_ind[0] = 1;
-
-    // Part 1 : sort kth firt elementZ
-    for (l=1; l<k; l++){
-      curr_row  = l * width;
-      curr_dist = p_dist[curr_row];
-      if (curr_dist<max_dist){
-        i=l-1;
-        for (int a=0; a<l-1; a++){
-          if (p_dist[a*width]>curr_dist){
-            i=a;
-            break;
-          }
-        }
-        for (j=l; j>i; j--){
-          p_dist[j*width] = p_dist[(j-1)*width];
-          p_ind[j*width]   = p_ind[(j-1)*width];
-        }
-        p_dist[i*width] = curr_dist;
-        p_ind[i*width]   = l+1;
-      } else {
-        p_ind[l*width] = l+1;
-      }
-      max_dist = p_dist[curr_row];
-    }
-
-    // Part 2 : insert element in the k-th first lines
-    max_row = (k-1)*width;
-    for (l=k; l<height; l++){
-      curr_dist = p_dist[l*width];
-      if (curr_dist<max_dist){
-        i=k-1;
-        for (int a=0; a<k-1; a++){
-          if (p_dist[a*width]>curr_dist){
-            i=a;
-            break;
-          }
-        }
-        for (j=k-1; j>i; j--){
-          p_dist[j*width] = p_dist[(j-1)*width];
-          p_ind[j*width]   = p_ind[(j-1)*width];
-        }
-        p_dist[i*width] = curr_dist;
-        p_ind[i*width]   = l+1;
-        max_dist             = p_dist[max_row];
-      }
-    }
-  }
-}
-
-
-/**
-  * Computes the square root of the first line (width-th first element)
-  * of the distance matrix.
-  *
-  * @param dist    distance matrix
-  * @param width   width of the distance matrix
-  * @param k       number of neighbors to consider
-  */
-__global__ void cuParallelSqrt(float *dist, int width, int k){
-    unsigned int xIndex = blockIdx.x * blockDim.x + threadIdx.x;
-    unsigned int yIndex = blockIdx.y * blockDim.y + threadIdx.y;
-  if (xIndex<width && yIndex<k)
-    dist[yIndex*width + xIndex] = sqrt(dist[yIndex*width + xIndex]);
-}
-
-
-//-----------------------------------------------------------------------------------------------//
-//                                   K-th NEAREST NEIGHBORS                                      //
-//-----------------------------------------------------------------------------------------------//
-
-/**
-  * K nearest neighbor algorithm
-  * - Initialize CUDA
-  * - Allocate device memory
-  * - Copy point sets (reference and query points) from host to device memory
-  * - Compute the distances + indexes to the k nearest neighbors for each query point
-  * - Copy distances from device to host memory
-  *
-  * @param ref_host      reference points ; pointer to linear matrix
-  * @param ref_nb        number of reference points ; width of the matrix
-  * @param query_host    query points ; pointer to linear matrix
-  * @param query_nb      number of query points ; width of the matrix
-  * @param dim           dimension of points ; height of the matrices
-  * @param k             number of neighbor to consider
-  * @param dist_host     distances to k nearest neighbors ; pointer to linear matrix
-  * @param dist_host     indexes of the k nearest neighbors ; pointer to linear matrix
-  *
-  */
-void knn_device(float* ref_dev, int ref_nb, float* query_dev, int query_nb,
-    int dim, int k, float* dist_dev, long* ind_dev, cudaStream_t stream){
-
-  // Grids and threads
-  dim3 g_16x16(query_nb/16, ref_nb/16, 1);
-  dim3 t_16x16(16, 16, 1);
-  if (query_nb%16 != 0) g_16x16.x += 1;
-  if (ref_nb  %16 != 0) g_16x16.y += 1;
-  //
-  dim3 g_256x1(query_nb/256, 1, 1);
-  dim3 t_256x1(256, 1, 1);
-  if (query_nb%256 != 0) g_256x1.x += 1;
-
-  dim3 g_k_16x16(query_nb/16, k/16, 1);
-  dim3 t_k_16x16(16, 16, 1);
-  if (query_nb%16 != 0) g_k_16x16.x += 1;
-  if (k  %16 != 0) g_k_16x16.y += 1;
-
-  // Kernel 1: Compute all the distances
-  cuComputeDistanceGlobal<<<g_16x16, t_16x16, 0, stream>>>(ref_dev, ref_nb, query_dev, query_nb, dim, dist_dev);
-
-  // Kernel 2: Sort each column
-  cuInsertionSort<<<g_256x1, t_256x1, 0, stream>>>(dist_dev, ind_dev, query_nb, ref_nb, k);
-
-  // Kernel 3: Compute square root of k first elements
-  // cuParallelSqrt<<<g_k_16x16,t_k_16x16, 0, stream>>>(dist_dev, query_nb, k);
-
-#if DEBUG
-  unsigned int  size_of_float = sizeof(float);
-  unsigned long size_of_long  = sizeof(long);
-
-  float* dist_host = new float[query_nb * k];
-  long*  idx_host  = new long[query_nb * k];
-
-  // Memory copy of output from device to host
-  cudaMemcpy(&dist_host[0], dist_dev,
-      query_nb * k *size_of_float, cudaMemcpyDeviceToHost);
-
-  cudaMemcpy(&idx_host[0], ind_dev,
-      query_nb * k * size_of_long, cudaMemcpyDeviceToHost);
-
-  int i = 0;
-  for(i = 0; i < 100; i++){
-    printf("IDX[%d]: %d\n", i, (int)idx_host[i]);
-  }
-#endif
-}
-
-
-
-
-
-

vision_detect/vision_detect/VisionDetect/net/knn/src/cuda/vision.h

@@ -1,8 +0,0 @@
-#pragma once
-#include <torch/extension.h>
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-
-void knn_device(float* ref_dev, int ref_width,
-    float* query_dev, int query_width,
-    int height, int k, float* dist_dev, long* ind_dev, cudaStream_t stream);

vision_detect/vision_detect/VisionDetect/net/knn/src/knn.h

@@ -1,81 +0,0 @@
-#pragma once
-#include "cpu/vision.h"
-
-#ifdef WITH_CUDA
-#include "cuda/vision.h"
-//#include <THC/THC.h>
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-//extern THCState *state;
-#endif
-
-
-
-int knn(at::Tensor& ref, at::Tensor& query, at::Tensor& idx)
-{
-
-    // NOTE: check dimensions
-    long batch, ref_nb, query_nb, dim, k;
-    batch = ref.size(0);
-    dim = ref.size(1);
-    k = idx.size(1);
-    ref_nb = ref.size(2);
-    query_nb = query.size(2);
-
-    float *ref_dev = ref.data<float>();
-    float *query_dev = query.data<float>();
-    long *idx_dev = idx.data<long>();
-
-
-
-
-  if (ref.type().is_cuda()) {
-#ifdef WITH_CUDA
-    // NOTE: raise error if not compiled with CUDA
-//    float *dist_dev = (float*)THCudaMalloc(state, ref_nb * query_nb * sizeof(float));
-    float *dist_dev = nullptr;
-    cudaError_t err = cudaMalloc((void**)&dist_dev, ref_nb * query_nb * sizeof(float));
-    if (err != cudaSuccess) {
-        fprintf(stderr, "cudaMalloc failed\n");
-        exit(1);
-    }
-
-    cudaFree(dist_dev);
-
-    for (int b = 0; b < batch; b++)
-    {
-    // knn_device(ref_dev + b * dim * ref_nb, ref_nb, query_dev + b * dim * query_nb, query_nb, dim, k,
-    //   dist_dev, idx_dev + b * k * query_nb, THCState_getCurrentStream(state));
-      knn_device(ref_dev + b * dim * ref_nb, ref_nb, query_dev + b * dim * query_nb, query_nb, dim, k,
-      dist_dev, idx_dev + b * k * query_nb, c10::cuda::getCurrentCUDAStream());
-    }
-//    THCudaFree(state, dist_dev);
-    cudaFree(dist_dev);
-    err = cudaGetLastError();
-    if (err != cudaSuccess)
-    {
-        printf("error in knn: %s\n", cudaGetErrorString(err));
-//        THError("aborting");
-        fprintf(stderr, "aborting\n");
-        exit(1);
-    }
-    return 1;
-#else
-    AT_ERROR("Not compiled with GPU support");
-#endif
-  }
-
-
-    float *dist_dev = (float*)malloc(ref_nb * query_nb * sizeof(float));
-    long *ind_buf = (long*)malloc(ref_nb * sizeof(long));
-    for (int b = 0; b < batch; b++) {
-    knn_cpu(ref_dev + b * dim * ref_nb, ref_nb, query_dev + b * dim * query_nb, query_nb, dim, k,
-      dist_dev, idx_dev + b * k * query_nb, ind_buf);
-    }
-
-    free(dist_dev);
-    free(ind_buf);
-
-    return 1;
-
-}

vision_detect/vision_detect/VisionDetect/net/knn/src/vision.cpp

@@ -1,5 +0,0 @@
-#include "knn.h"
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("knn", &knn, "k-nearest neighbors");
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/include/ball_query.h

@@ -1,11 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#pragma once
-#include <torch/extension.h>
-
-
-at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius,
-                      const int nsample);

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/include/cuda_utils.h

@@ -1,46 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#ifndef _CUDA_UTILS_H
-#define _CUDA_UTILS_H
-
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <cmath>
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-#include <vector>
-
-#define TOTAL_THREADS 512
-
-inline int opt_n_threads(int work_size) {
-  const int pow_2 = std::log(static_cast<double>(work_size)) / std::log(2.0);
-
-  return max(min(1 << pow_2, TOTAL_THREADS), 1);
-}
-
-inline dim3 opt_block_config(int x, int y) {
-  const int x_threads = opt_n_threads(x);
-  const int y_threads =
-      max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1);
-  dim3 block_config(x_threads, y_threads, 1);
-
-  return block_config;
-}
-
-#define CUDA_CHECK_ERRORS()                                           \
-  do {                                                                \
-    cudaError_t err = cudaGetLastError();                             \
-    if (cudaSuccess != err) {                                         \
-      fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n",  \
-              cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \
-              __FILE__);                                              \
-      exit(-1);                                                       \
-    }                                                                 \
-  } while (0)
-
-#endif

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/include/cylinder_query.h

@@ -1,7 +0,0 @@
-// Author: chenxi-wang
-
-#pragma once
-#include <torch/extension.h>
-
-at::Tensor cylinder_query(at::Tensor new_xyz, at::Tensor xyz, at::Tensor rot, const float radius, const float hmin, const float hmax,
-                      const int nsample);

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/include/group_points.h

@@ -1,10 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#pragma once
-#include <torch/extension.h>
-
-at::Tensor group_points(at::Tensor points, at::Tensor idx);
-at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n);

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/include/interpolate.h

@@ -1,15 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#pragma once
-
-#include <torch/extension.h>
-#include <vector>
-
-std::vector<at::Tensor> three_nn(at::Tensor unknowns, at::Tensor knows);
-at::Tensor three_interpolate(at::Tensor points, at::Tensor idx,
-                             at::Tensor weight);
-at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx,
-                                  at::Tensor weight, const int m);

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/include/sampling.h

@@ -1,11 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#pragma once
-#include <torch/extension.h>
-
-at::Tensor gather_points(at::Tensor points, at::Tensor idx);
-at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n);
-at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples);

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/include/utils.h

@@ -1,32 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#pragma once
-#ifdef WITH_CUDA
-#include <ATen/cuda/CUDAContext.h>
-#endif
-#include <torch/extension.h>
-
-#define CHECK_CUDA(x)                                          \
-  do {                                                         \
-    TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor"); \
-  } while (0)
-
-#define CHECK_CONTIGUOUS(x)                                         \
-  do {                                                              \
-    TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \
-  } while (0)
-
-#define CHECK_IS_INT(x)                              \
-  do {                                               \
-    TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, \
-             #x " must be an int tensor");           \
-  } while (0)
-
-#define CHECK_IS_FLOAT(x)                              \
-  do {                                                 \
-    TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, \
-             #x " must be a float tensor");            \
-  } while (0)

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/ball_query.cpp

@@ -1,113 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-#include<iostream>
-#include "ball_query.h"
-#include "utils.h"
-
-
-#ifdef WITH_CUDA
-void query_ball_point_kernel_wrapper(
-  int b, int n, int m, float radius,
-  int nsample, 
-  const float *new_xyz,
-  const float *xyz, 
-  int *idx
-);
-#endif
-
-void query_ball_point_cpu(
-  int b, int n, int m, float radius,
-  int nsample,
-  const float *new_xyz,
-  const float *xyz, 
-  int *idx
-);
-
-#ifdef DEBUG
-void ball_query_debug(const torch::Tensor& new_xyz, const torch::Tensor& xyz, 
-                      float radius, int nsample) {
-    // 优化后的打印代码，格式更清晰，便于调试
-    std::cout << "===== ball_query.cpp 调试信息 =====" << std::endl;
-    // 打印new_xyz的尺寸（维度）
-    std::cout << "new_xyz.size(): (";
-    for (auto s : new_xyz.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" <<std::endl;
-    
-    // 打印xyz的尺寸（维度）
-    std::cout << "xyz.size():    (";
-    for (auto s : xyz.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-
-    // 打印半径和采样数
-    std::cout << "radius:        " << radius << std::endl;
-    std::cout << "nsample:       " << nsample << std::endl;
-    std::cout << "====================================" << std::endl;
-}
-#endif
-
-
-at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius,
-                      const int nsample) {
-  CHECK_CONTIGUOUS(new_xyz);
-  CHECK_CONTIGUOUS(xyz);
-  CHECK_IS_FLOAT(new_xyz);
-  CHECK_IS_FLOAT(xyz);
-
-  #ifdef WITH_CUDA
-  if (new_xyz.is_cuda()) {
-    CHECK_CUDA(xyz);
-    // 新增1：CUDA张量设备端连续性强校验（解决内存布局碎片化问题）
-    TORCH_CHECK(new_xyz.is_contiguous(), "new_xyz must be contiguous on CUDA device");
-    TORCH_CHECK(xyz.is_contiguous(), "xyz must be contiguous on CUDA device");
-    // 新增2：CUDA设备同步，确保张量内存分配完成（解决指针无效问题）
-    cudaDeviceSynchronize();
-  } 
-  #endif
-
-  #ifdef DEBUG
-  ball_query_debug(new_xyz, xyz, radius, nsample);
-  #endif
-
-  at::Tensor idx =
-      torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample},
-                   at::device(new_xyz.device()).dtype(at::ScalarType::Int));
-  
-
-  #ifdef WITH_CUDA
-  if (new_xyz.is_cuda()) {
-    query_ball_point_kernel_wrapper(
-      xyz.size(0), 
-      xyz.size(1), 
-      new_xyz.size(1),                      
-      radius, nsample,  
-      new_xyz.contiguous().data_ptr<float>(),
-      xyz.contiguous().data_ptr<float>(), 
-      idx.contiguous().data_ptr<int>()
-    );
-
-  } else {
-  #endif
-    // TORCH_CHECK(false, "CPU not supported");
-    query_ball_point_cpu(
-        xyz.size(0), 
-        xyz.size(1), 
-        new_xyz.size(1), 
-        radius, 
-        nsample,
-        new_xyz.contiguous().data_ptr<float>(), 
-        xyz.contiguous().data_ptr<float>(), 
-        idx.contiguous().data_ptr<int>()
-      );
-    
-  #ifdef WITH_CUDA
-  }
-  #endif
-  
-  return idx;
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/ball_query_cpu.cpp

@@ -1,44 +0,0 @@
-
-
-void query_ball_point_cpu(
-    int b, int n, int m, float radius,
-    int nsample,
-    const float* new_xyz,
-    const float* xyz,
-    int* idx
-) {
-    float radius2 = radius * radius;
-
-    for (int batch_index = 0; batch_index < b; ++batch_index) {
-        const float* xyz_b = xyz + batch_index * n * 3;
-        const float* new_xyz_b = new_xyz + batch_index * m * 3;
-        int* idx_b = idx + batch_index * m * nsample;
-
-        for (int j = 0; j < m; ++j) {
-            float nx = new_xyz_b[j * 3 + 0];
-            float ny = new_xyz_b[j * 3 + 1];
-            float nz = new_xyz_b[j * 3 + 2];
-
-            for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) {
-                float x = xyz_b[k * 3 + 0];
-                float y = xyz_b[k * 3 + 1];
-                float z = xyz_b[k * 3 + 2];
-
-                float d2 =
-                    (nx - x) * (nx - x) +
-                    (ny - y) * (ny - y) +
-                    (nz - z) * (nz - z);
-
-                if (d2 < radius2) {
-                    if (cnt == 0) {
-                        for (int l = 0; l < nsample; ++l) {
-                            idx_b[j * nsample + l] = k;
-                        }
-                    }
-                    idx_b[j * nsample + cnt] = k;
-                    ++cnt;
-                }
-            }
-        }
-    }
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/ball_query_gpu.cu

@@ -1,59 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#include <math.h>
-#include <stdio.h>
-#include <stdlib.h>
-
-#include "cuda_utils.h"
-
-// input: new_xyz(b, m, 3) xyz(b, n, 3)
-// output: idx(b, m, nsample)
-__global__ void query_ball_point_kernel(int b, int n, int m, float radius,
-                                        int nsample,
-                                        const float *__restrict__ new_xyz,
-                                        const float *__restrict__ xyz,
-                                        int *__restrict__ idx) {
-  int batch_index = blockIdx.x;
-  xyz += batch_index * n * 3;
-  new_xyz += batch_index * m * 3;
-  idx += m * nsample * batch_index;
-
-  int index = threadIdx.x;
-  int stride = blockDim.x;
-
-  float radius2 = radius * radius;
-  for (int j = index; j < m; j += stride) {
-    float new_x = new_xyz[j * 3 + 0];
-    float new_y = new_xyz[j * 3 + 1];
-    float new_z = new_xyz[j * 3 + 2];
-    for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) {
-      float x = xyz[k * 3 + 0];
-      float y = xyz[k * 3 + 1];
-      float z = xyz[k * 3 + 2];
-      float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) +
-                 (new_z - z) * (new_z - z);
-      if (d2 < radius2) {
-        if (cnt == 0) {
-          for (int l = 0; l < nsample; ++l) {
-            idx[j * nsample + l] = k;
-          }
-        }
-        idx[j * nsample + cnt] = k;
-        ++cnt;
-      }
-    }
-  }
-}
-
-void query_ball_point_kernel_wrapper(int b, int n, int m, float radius,
-                                     int nsample, const float *new_xyz,
-                                     const float *xyz, int *idx) {
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  query_ball_point_kernel<<<b, opt_n_threads(m), 0, stream>>>(
-      b, n, m, radius, nsample, new_xyz, xyz, idx);
-
-  CUDA_CHECK_ERRORS();
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/bindings.cpp

@@ -1,27 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#include "ball_query.h"
-#include "group_points.h"
-#include "interpolate.h"
-#include "sampling.h"
-#include "cylinder_query.h"
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("gather_points", &gather_points);
-  m.def("gather_points_grad", &gather_points_grad);
-  m.def("furthest_point_sampling", &furthest_point_sampling);
-
-  m.def("three_nn", &three_nn);
-  m.def("three_interpolate", &three_interpolate);
-  m.def("three_interpolate_grad", &three_interpolate_grad);
-
-  m.def("ball_query", &ball_query);
-
-  m.def("group_points", &group_points);
-  m.def("group_points_grad", &group_points_grad);
-
-  m.def("cylinder_query", &cylinder_query);
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/cylinder_query.cpp

@@ -1,95 +0,0 @@
-// Author: chenxi-wang
-
-#include "cylinder_query.h"
-#include "utils.h"
-
-
-#ifdef WITH_CUDA
-void query_cylinder_point_kernel_wrapper(int b, int n, int m, float radius, float hmin, float hmax,
-                                     int nsample, const float *new_xyz,
-                                     const float *xyz, const float *rot, int *idx);
-#endif
-
-
-void query_cylinder_point_cpu(
-    int b, int n, int m, float radius, float hmin, float hmax,
-    int nsample, const float *new_xyz,
-    const float *xyz, const float *rot, int *idx);
-
-#ifdef DEBUG
-void cylinder_query_debug(at::Tensor new_xyz, at::Tensor xyz, at::Tensor rot, 
-                      const float radius, const float hmin, const float hmax,
-                      const int nsample) {
-    // 优化后的打印代码，格式更清晰，便于调试
-    std::cout << "===== cylinder_query.cpp 调试信息 =====" << std::endl;
-    // 打印new_xyz的尺寸（维度）
-    std::cout << "new_xyz.size(): (";
-    for (auto s : new_xyz.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" <<std::endl;
-    
-    // 打印xyz的尺寸（维度）
-    std::cout << "xyz.size():    (";
-    for (auto s : xyz.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-
-    // 打印rot的尺寸（维度）
-    std::cout << "rot.size():    (";
-    for (auto s : rot.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-    
-    // 打印半径和采样数
-    std::cout << "radius:        " << radius << std::endl;
-    std::cout << "hmin:          " << hmin << std::endl;
-    std::cout << "hmax:          " << hmax << std::endl;
-    std::cout << "nsample:       " << nsample << std::endl;
-    std::cout << "====================================" << std::endl;
-}
-#endif
-
-at::Tensor cylinder_query(at::Tensor new_xyz, at::Tensor xyz, at::Tensor rot, const float radius, const float hmin, const float hmax,
-                      const int nsample) {
-  CHECK_CONTIGUOUS(new_xyz);
-  CHECK_CONTIGUOUS(xyz);
-  CHECK_CONTIGUOUS(rot);
-  CHECK_IS_FLOAT(new_xyz);
-  CHECK_IS_FLOAT(xyz);
-  CHECK_IS_FLOAT(rot);
-
-  #ifdef WITH_CUDA
-  if (new_xyz.is_cuda()) {
-    CHECK_CUDA(xyz);
-    CHECK_CUDA(rot);
-  }
-  #endif
-
-  #ifdef DEBUG
-  cylinder_query_debug(new_xyz, xyz, rot, radius, hmin, hmax, nsample);
-  #endif
-
-  at::Tensor idx =
-      torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample},
-                   at::device(new_xyz.device()).dtype(at::ScalarType::Int));
-
-  #ifdef WITH_CUDA
-  if (new_xyz.is_cuda()) {
-    query_cylinder_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1),
-                                    radius, hmin, hmax, nsample, new_xyz.data_ptr<float>(),
-                                    xyz.data_ptr<float>(), rot.data_ptr<float>(), idx.data_ptr<int>());
-  } else {
-  #endif
-    // TORCH_CHECK(false, "CPU not supported");
-    query_cylinder_point_cpu(xyz.size(0), xyz.size(1), new_xyz.size(1),
-                             radius, hmin, hmax, nsample, new_xyz.data_ptr<float>(),
-                             xyz.data_ptr<float>(), rot.data_ptr<float>(), idx.data_ptr<int>());
-  #ifdef WITH_CUDA
-  }
-  #endif
-
-  return idx;
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/cylinder_query_cpu.cpp

@@ -1,56 +0,0 @@
-void query_cylinder_point_cpu(
-    int b, int n, int m, 
-    float radius, 
-    float hmin, float hmax,
-    int nsample, 
-    const float *new_xyz,
-    const float *xyz, 
-    const float *rot, 
-    int *idx
-) {
-    float radius2 = radius * radius;
-
-    for (int batch_index = 0; batch_index < b; ++batch_index) {
-        const float *xyz_b = xyz + batch_index * n * 3;
-        const float *new_xyz_b = new_xyz + batch_index * m * 3;
-        const float *rot_b = rot + batch_index * m * 9;
-        int *idx_b = idx + batch_index * m * nsample;
-
-        for (int j = 0; j < m; ++j) {
-            float nx = new_xyz_b[j * 3 + 0];
-            float ny = new_xyz_b[j * 3 + 1];
-            float nz = new_xyz_b[j * 3 + 2];
-
-            float r0 = rot_b[j * 9 + 0];
-            float r1 = rot_b[j * 9 + 1];
-            float r2 = rot_b[j * 9 + 2];
-            float r3 = rot_b[j * 9 + 3];
-            float r4 = rot_b[j * 9 + 4];
-            float r5 = rot_b[j * 9 + 5];
-            float r6 = rot_b[j * 9 + 6];
-            float r7 = rot_b[j * 9 + 7];
-            float r8 = rot_b[j * 9 + 8];
-
-            for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) {
-                float x = xyz_b[k * 3 + 0] - nx;
-                float y = xyz_b[k * 3 + 1] - ny;
-                float z = xyz_b[k * 3 + 2] - nz;
-
-                float x_rot = r0 * x + r3 * y + r6 * z;
-                float y_rot = r1 * x + r4 * y + r7 * z;
-                float z_rot = r2 * x + r5 * y + r8 * z;
-
-                float d2 = y_rot * y_rot + z_rot * z_rot;
-                if (d2 < radius2 && x_rot > hmin && x_rot < hmax) {
-                    if (cnt == 0) {
-                        for (int l = 0; l < nsample; ++l) {
-                            idx_b[j * nsample + l] = k;
-                        }
-                    }
-                    idx_b[j * nsample + cnt] = k;
-                    ++cnt;
-                }
-            }
-        }
-    }
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/cylinder_query_gpu.cu

@@ -1,67 +0,0 @@
-// Author: chenxi-wang
-
-#include <math.h>
-#include <stdio.h>
-#include <stdlib.h>
-
-#include "cuda_utils.h"
-
-__global__ void query_cylinder_point_kernel(int b, int n, int m, float radius, float hmin, float hmax,
-                                        int nsample,
-                                        const float *__restrict__ new_xyz,
-                                        const float *__restrict__ xyz,
-                                        const float *__restrict__ rot,
-                                        int *__restrict__ idx) {
-  int batch_index = blockIdx.x;
-  xyz += batch_index * n * 3;
-  new_xyz += batch_index * m * 3;
-  rot += batch_index * m * 9;
-  idx += m * nsample * batch_index;
-
-  int index = threadIdx.x;
-  int stride = blockDim.x;
-
-  float radius2 = radius * radius;
-  for (int j = index; j < m; j += stride) {
-    float new_x = new_xyz[j * 3 + 0];
-    float new_y = new_xyz[j * 3 + 1];
-    float new_z = new_xyz[j * 3 + 2];
-    float r0 = rot[j * 9 + 0];
-    float r1 = rot[j * 9 + 1];
-    float r2 = rot[j * 9 + 2];
-    float r3 = rot[j * 9 + 3];
-    float r4 = rot[j * 9 + 4];
-    float r5 = rot[j * 9 + 5];
-    float r6 = rot[j * 9 + 6];
-    float r7 = rot[j * 9 + 7];
-    float r8 = rot[j * 9 + 8];
-    for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) {
-      float x = xyz[k * 3 + 0] - new_x;
-      float y = xyz[k * 3 + 1] - new_y;
-      float z = xyz[k * 3 + 2] - new_z;
-      float x_rot = r0 * x + r3 * y + r6 * z;
-      float y_rot = r1 * x + r4 * y + r7 * z;
-      float z_rot = r2 * x + r5 * y + r8 * z;
-      float d2 = y_rot * y_rot + z_rot * z_rot;
-      if (d2 < radius2 && x_rot > hmin && x_rot < hmax) {
-        if (cnt == 0) {
-          for (int l = 0; l < nsample; ++l) {
-            idx[j * nsample + l] = k;
-          }
-        }
-        idx[j * nsample + cnt] = k;
-        ++cnt;
-      }
-    }
-  }
-}
-
-void query_cylinder_point_kernel_wrapper(int b, int n, int m, float radius, float hmin, float hmax,
-                                     int nsample, const float *new_xyz,
-                                     const float *xyz, const float *rot, int *idx) {
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  query_cylinder_point_kernel<<<b, opt_n_threads(m), 0, stream>>>(
-      b, n, m, radius, hmin, hmax, nsample, new_xyz, xyz, rot, idx);
-
-  CUDA_CHECK_ERRORS();
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/group_points.cpp

@@ -1,153 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#include "group_points.h"
-#include "utils.h"
-
-
-#ifdef WITH_CUDA
-void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample,
-                                 const float *points, const int *idx,
-                                 float *out);
-
-void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
-                                      int nsample, const float *grad_out,
-                                      const int *idx, float *grad_points);
-#endif
-
-
-void group_points_cpu(int b, int c, int n, int npoints, int nsample,
-                      const float *points, const int *idx,
-                      float *out);
-
-void group_points_grad_cpu(int b, int c, int n, int npoints,
-                          int nsample, const float *grad_out,
-                          const int *idx, float *grad_points);
-
-#ifdef DEBUG
-void group_points_debug(at::Tensor points, at::Tensor idx) {
-    // 优化后的打印代码，格式更清晰，便于调试
-    std::cout << "===== group_points.cpp 调试信息 =====" << std::endl;
-    std::cout << "=====  group points 函数调试信息 =====" << std::endl;
-    // 打印new_xyz的尺寸（维度）
-    std::cout << "points.size(): (";
-    for (auto s : points.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" <<std::endl;
-    
-    // 打印xyz的尺寸（维度）
-    std::cout << "idx.size():    (";
-    for (auto s : idx.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-
-    std::cout << "====================================" << std::endl;
-}
-
-void group_points_grad_debug(at::Tensor grad_out, at::Tensor idx, const int n) {
-    // 优化后的打印代码，格式更清晰，便于调试
-    std::cout << "===== group_points.cpp 调试信息 =====" << std::endl;
-    std::cout << "===== group_points_grad 函数调试信息 =====" << std::endl;
-    // 打印new_xyz的尺寸（维度）
-    std::cout << "grad_out.size(): (";
-    for (auto s : grad_out.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" <<std::endl;
-    
-    // 打印xyz的尺寸（维度）
-    std::cout << "idx.size():    (";
-    for (auto s : idx.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-    
-    std::cout << "n:             " << n << std::endl;
-    std::cout << "====================================" << std::endl;
-}
-#endif
-
-                          
-
-at::Tensor group_points(at::Tensor points, at::Tensor idx) {
-  CHECK_CONTIGUOUS(points);
-  CHECK_CONTIGUOUS(idx);
-  CHECK_IS_FLOAT(points);
-  CHECK_IS_INT(idx);
-
-  #ifdef WITH_CUDA
-  if (points.is_cuda()) {
-    CHECK_CUDA(idx);
-  }
-  #endif
-
-  #ifdef DEBUG
-  group_points_debug(points, idx);
-  #endif
-
-  at::Tensor output =
-      torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)},
-                   at::device(points.device()).dtype(at::ScalarType::Float));
-
-  #ifdef WITH_CUDA
-  if (points.is_cuda()) {
-    group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2),
-                                idx.size(1), idx.size(2), points.data_ptr<float>(),
-                                idx.data_ptr<int>(), output.data_ptr<float>());
-  } else {
-  #endif
-    // TORCH_CHECK(false, "CPU not supported");
-    group_points_cpu(points.size(0), points.size(1), points.size(2),
-                     idx.size(1), idx.size(2), points.data_ptr<float>(),
-                     idx.data_ptr<int>(), output.data_ptr<float>());
-
-  #ifdef WITH_CUDA
-  }
-  #endif
-
-  return output;
-}
-
-at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) {
-  CHECK_CONTIGUOUS(grad_out);
-  CHECK_CONTIGUOUS(idx);
-  CHECK_IS_FLOAT(grad_out);
-  CHECK_IS_INT(idx);
-
-  #ifdef WITH_CUDA
-  if (grad_out.is_cuda()) {
-    CHECK_CUDA(idx);
-  }
-  #endif
-
-  #ifdef DEBUG
-  group_points_grad_debug(grad_out, idx, n);
-  #endif
-
-  at::Tensor output =
-      torch::zeros({grad_out.size(0), grad_out.size(1), n},
-                   at::device(grad_out.device()).dtype(at::ScalarType::Float));
-
-  #ifdef WITH_CUDA
-  if (grad_out.is_cuda()) {
-    group_points_grad_kernel_wrapper(
-        grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2),
-        grad_out.data_ptr<float>(), idx.data_ptr<int>(), output.data_ptr<float>());
-  } else {
-  #endif
-
-    // TORCH_CHECK(false, "CPU not supported");
-    group_points_grad_cpu(
-        grad_out.size(0), grad_out.size(1), n,idx.size(1), idx.size(2),
-        grad_out.data_ptr<float>(), idx.data_ptr<int>(), output.data_ptr<float>());
-
-  #ifdef WITH_CUDA
-  } 
-  #endif 
-
-  return output;
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/group_points_cpu.cpp

@@ -1,54 +0,0 @@
-void group_points_cpu(
-    int b, int c, int n, 
-    int npoints, 
-    int nsample,            
-    const float *points, 
-    const int *idx,        
-    float *out   
-) {
-    for (int batch = 0; batch < b; ++batch) {
-        const float *points_b = points + batch * n * c;
-        const int *idx_b = idx + batch * npoints * nsample;
-        float *out_b = out + batch * npoints * nsample * c;
-
-        for (int i = 0; i < c * npoints; ++i) {
-            const int l = i / npoints;
-            const int j = i % npoints;
-            
-            for (int k = 0; k < nsample; ++k) {
-                int ii = idx_b[j * nsample + k];
-                out_b[(l * npoints + j) * nsample + k] = points_b[l * n + ii];
-            }
-        }
-    }
-}
-
-void group_points_grad_cpu(
-    int b, int c, int n, 
-    int npoints,      
-    int nsample, 
-    const float *grad_out,        
-    const int *idx, 
-    float *grad_points
-) {
-    for (int batch_index = 0; batch_index < b; ++batch_index) {
-        const float *grad_out_b = grad_out + batch_index * npoints * nsample * c;
-        const int *idx_b = idx + batch_index * npoints * nsample;
-        float *grad_points_b = grad_points + batch_index * n * c;
-
-        for (int i = 0; i < n * c; ++i) {
-            grad_points_b[i] = 0.0f;
-        }
-
-        for (int i = 0; i < c * npoints; ++i) {
-            const int l = i / npoints;
-            const int j = i % npoints;
-
-            for (int k = 0; k < nsample; ++k) {
-                int ii = idx_b[j * nsample + k];
-                grad_points_b[l * n + ii] += grad_out_b[(l * npoints + j) * nsample + k];
-            }
-        }
-                
-    }
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/group_points_gpu.cu

@@ -1,84 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#include <stdio.h>
-#include <stdlib.h>
-
-#include "cuda_utils.h"
-
-// input: points(b, c, n) idx(b, npoints, nsample)
-// output: out(b, c, npoints, nsample)
-__global__ void group_points_kernel(int b, int c, int n, int npoints,
-                                    int nsample,
-                                    const float *__restrict__ points,
-                                    const int *__restrict__ idx,
-                                    float *__restrict__ out) {
-  int batch_index = blockIdx.x;
-  points += batch_index * n * c;
-  idx += batch_index * npoints * nsample;
-  out += batch_index * npoints * nsample * c;
-
-  const int index = threadIdx.y * blockDim.x + threadIdx.x;
-  const int stride = blockDim.y * blockDim.x;
-  for (int i = index; i < c * npoints; i += stride) {
-    const int l = i / npoints;
-    const int j = i % npoints;
-    for (int k = 0; k < nsample; ++k) {
-      int ii = idx[j * nsample + k];
-      if (ii < 0 || ii >= n) {
-        out[(l * npoints + j) * nsample + k] = 0;
-      } else {
-        out[(l * npoints + j) * nsample + k] = points[l * n + ii];
-      }
-    }
-  }
-}
-
-void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample,
-                                 const float *points, const int *idx,
-                                 float *out) {
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  group_points_kernel<<<b, opt_block_config(npoints, c), 0, stream>>>(
-      b, c, n, npoints, nsample, points, idx, out);
-
-  CUDA_CHECK_ERRORS();
-}
-
-// input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample)
-// output: grad_points(b, c, n)
-__global__ void group_points_grad_kernel(int b, int c, int n, int npoints,
-                                         int nsample,
-                                         const float *__restrict__ grad_out,
-                                         const int *__restrict__ idx,
-                                         float *__restrict__ grad_points) {
-  int batch_index = blockIdx.x;
-  grad_out += batch_index * npoints * nsample * c;
-  idx += batch_index * npoints * nsample;
-  grad_points += batch_index * n * c;
-
-  const int index = threadIdx.y * blockDim.x + threadIdx.x;
-  const int stride = blockDim.y * blockDim.x;
-  for (int i = index; i < c * npoints; i += stride) {
-    const int l = i / npoints;
-    const int j = i % npoints;
-    for (int k = 0; k < nsample; ++k) {
-      int ii = idx[j * nsample + k];
-      atomicAdd(grad_points + l * n + ii,
-                grad_out[(l * npoints + j) * nsample + k]);
-    }
-  }
-}
-
-void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
-                                      int nsample, const float *grad_out,
-                                      const int *idx, float *grad_points) {
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  group_points_grad_kernel<<<b, opt_block_config(npoints, c), 0, stream>>>(
-      b, c, n, npoints, nsample, grad_out, idx, grad_points);
-
-  CUDA_CHECK_ERRORS();
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/interpolate.cpp

@@ -1,241 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#include "interpolate.h"
-#include "utils.h"
-
-
-#ifdef WITH_CUDA
-void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown,
-                             const float *known, float *dist2, int *idx);
-void three_interpolate_kernel_wrapper(int b, int c, int m, int n,
-                                      const float *points, const int *idx,
-                                      const float *weight, float *out);
-void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m,
-                                           const float *grad_out,
-                                           const int *idx, const float *weight,
-                                           float *grad_points);
-#endif
-
-
-void three_nn_cpu(int b, int n, int m, const float *unknown,
-                             const float *known, float *dist2, int *idx);
-void three_interpolate_cpu(int b, int c, int m, int n,
-                                      const float *points, const int *idx,
-                                      const float *weight, float *out);
-void three_interpolate_grad_cpu(int b, int c, int n, int m,
-                                           const float *grad_out,
-                                           const int *idx, const float *weight,
-                                           float *grad_points);
-
-#ifdef DEBUG
-void three_nn_debug(at::Tensor unknowns, at::Tensor knows) {
-    std::cout << "===== interpolate.cpp 调试信息 =====" << std::endl;
-    std::cout << "=========== three_nn ===========" << std::endl;
-
-    std::cout << "unknowns.size(): (";
-    for (auto s : unknowns.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" <<std::endl;
-    
-    std::cout << "knows.size():    (";
-    for (auto s : knows.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-
-    std::cout << "====================================" << std::endl;
-}
-
-void three_interpolate_debug(at::Tensor points, at::Tensor idx,
-                             at::Tensor weight) {
-    std::cout << "===== interpolate.cpp 调试信息 =====" << std::endl;
-    std::cout << "=========== three_interpolate ===========" << std::endl;
-
-    std::cout << "points.size(): (";
-    for (auto s : points.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" <<std::endl;
-    
-    std::cout << "idx.size():    (";
-    for (auto s : idx.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-
-    std::cout << "weight.size():    (";
-    for (auto s : weight.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-
-    std::cout << "====================================" << std::endl;
-}
-
-void three_interpolate_grad_debug(at::Tensor grad_out, at::Tensor idx,
-                                  at::Tensor weight, const int m) {
-    std::cout << "===== interpolate.cpp 调试信息 =====" << std::endl;
-    std::cout << "=========== three_interpolate_grad ===========" << std::endl;
-
-    std::cout << "grad_out.size(): (";
-    for (auto s : grad_out.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" <<std::endl;
-    
-    std::cout << "idx.size():    (";
-    for (auto s : idx.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-
-    std::cout << "weight.size():    (";
-    for (auto s : weight.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-    
-    // 打印半径和采样数
-    std::cout << "m:             " << m << std::endl;
-    std::cout << "====================================" << std::endl;
-}
-#endif
-
-
-std::vector<at::Tensor> three_nn(at::Tensor unknowns, at::Tensor knows) {
-  CHECK_CONTIGUOUS(unknowns);
-  CHECK_CONTIGUOUS(knows);
-  CHECK_IS_FLOAT(unknowns);
-  CHECK_IS_FLOAT(knows);
-
-  #ifdef WITH_CUDA
-  if (unknowns.is_cuda()) {
-    CHECK_CUDA(knows);
-  }
-  #endif
-
-  #ifdef DEBUG
-  three_nn_debug(unknowns, knows);
-  #endif
-
-  at::Tensor idx =
-      torch::zeros({unknowns.size(0), unknowns.size(1), 3},
-                   at::device(unknowns.device()).dtype(at::ScalarType::Int));
-  at::Tensor dist2 =
-      torch::zeros({unknowns.size(0), unknowns.size(1), 3},
-                   at::device(unknowns.device()).dtype(at::ScalarType::Float));
-
-  #ifdef WITH_CUDA
-  if (unknowns.is_cuda()) {
-    three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1),
-                            unknowns.data_ptr<float>(), knows.data_ptr<float>(),
-                            dist2.data_ptr<float>(), idx.data_ptr<int>());
-  } else {
-  #endif
-  
-    // TORCH_CHECK(false, "CPU not supported");
-    three_nn_cpu(unknowns.size(0), unknowns.size(1), knows.size(1),
-                 unknowns.data_ptr<float>(), knows.data_ptr<float>(),
-                 dist2.data_ptr<float>(), idx.data_ptr<int>());
-  #ifdef WITH_CUDA
-  }
-  #endif
-
-  return {dist2, idx};
-}
-
-at::Tensor three_interpolate(at::Tensor points, at::Tensor idx,
-                             at::Tensor weight) {
-  CHECK_CONTIGUOUS(points);
-  CHECK_CONTIGUOUS(idx);
-  CHECK_CONTIGUOUS(weight);
-  CHECK_IS_FLOAT(points);
-  CHECK_IS_INT(idx);
-  CHECK_IS_FLOAT(weight);
-
-
-  #ifdef WITH_CUDA
-  if (points.is_cuda()) {
-    CHECK_CUDA(idx);
-    CHECK_CUDA(weight);
-  }
-  #endif
-
-  #ifdef DEBUG
-  three_interpolate_debug(points, idx, weight);
-  #endif
-
-  at::Tensor output =
-      torch::zeros({points.size(0), points.size(1), idx.size(1)},
-                   at::device(points.device()).dtype(at::ScalarType::Float));
-
-  #ifdef WITH_CUDA
-  if (points.is_cuda()) {
-    three_interpolate_kernel_wrapper(
-        points.size(0), points.size(1), points.size(2), idx.size(1),
-        points.data_ptr<float>(), idx.data_ptr<int>(), weight.data_ptr<float>(),
-        output.data_ptr<float>());
-  } else {
-  #endif
-
-    // TORCH_CHECK(false, "CPU not supported");
-    three_interpolate_cpu(
-        points.size(0), points.size(1), points.size(2), idx.size(1),
-        points.data_ptr<float>(), idx.data_ptr<int>(), weight.data_ptr<float>(),
-        output.data_ptr<float>());
-  #ifdef WITH_CUDA
-  }
-  #endif
-
-  return output;
-}
-
-at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx,
-                                  at::Tensor weight, const int m) {
-  CHECK_CONTIGUOUS(grad_out);
-  CHECK_CONTIGUOUS(idx);
-  CHECK_CONTIGUOUS(weight);
-  CHECK_IS_FLOAT(grad_out);
-  CHECK_IS_INT(idx);
-  CHECK_IS_FLOAT(weight);
-
-  #ifdef WITH_CUDA
-  if (grad_out.is_cuda()) {
-    CHECK_CUDA(idx);
-    CHECK_CUDA(weight);
-  }
-  #endif
-
-  #ifdef DEBUG
-  three_interpolate_grad_debug(grad_out, idx, weight, m);
-  #endif
-
-  at::Tensor output =
-      torch::zeros({grad_out.size(0), grad_out.size(1), m},
-                   at::device(grad_out.device()).dtype(at::ScalarType::Float));
-
-  #ifdef WITH_CUDA
-  if (grad_out.is_cuda()) {
-    three_interpolate_grad_kernel_wrapper(
-        grad_out.size(0), grad_out.size(1), grad_out.size(2), m,
-        grad_out.data_ptr<float>(), idx.data_ptr<int>(), weight.data_ptr<float>(),
-        output.data_ptr<float>());
-  } else {
-  #endif
-
-    // TORCH_CHECK(false, "CPU not supported");
-    three_interpolate_grad_cpu(
-        grad_out.size(0), grad_out.size(1), grad_out.size(2), m,
-        grad_out.data_ptr<float>(), idx.data_ptr<int>(), weight.data_ptr<float>(),
-        output.data_ptr<float>());
-        
-  #ifdef WITH_CUDA
-  }
-  #endif
-
-  return output;
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/interpolate_cpu.cpp

@@ -1,123 +0,0 @@
-void three_nn_cpu(
-    int b, int n, int m, 
-    const float *unknown,           
-    const float *known, 
-    float *dist2, 
-    int *idx
-) {
-    for (int batch_index = 0; batch_index < b; ++batch_index) {
-        const float *unknown_batch = unknown + batch_index * n * 3;
-        const float *known_batch = known + batch_index * m * 3;
-        float *dist2_batch = dist2 + batch_index * n * 3;
-        int *idx_batch = idx + batch_index * n * 3;
-
-        for (int j = 0; j < n; ++j) {
-            float ux = unknown_batch[j * 3 + 0];
-            float uy = unknown_batch[j * 3 + 1];
-            float uz = unknown_batch[j * 3 + 2];
-
-            double best1 = 1e40, best2 = 1e40, best3 = 1e40;
-            int besti1 = 0, besti2 = 0, besti3 = 0;
-
-            for (int k = 0; k < m; ++k) {
-                float x = known_batch[k * 3 + 0];
-                float y = known_batch[k * 3 + 1];
-                float z = known_batch[k * 3 + 2];
-                float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z);
-                if (d < best1) {
-                    best3 = best2;
-                    besti3 = besti2;
-                    best2 = best1;
-                    besti2 = besti1;
-                    best1 = d;
-                    besti1 = k;
-                } else if (d < best2) {
-                    best3 = best2;
-                    besti3 = besti2;
-                    best2 = d;
-                    besti2 = k;
-                } else if (d < best3) {
-                    best3 = d;
-                    besti3 = k;
-                }
-            }
-            dist2_batch[j * 3 + 0] = best1;
-            dist2_batch[j * 3 + 1] = best2;
-            dist2_batch[j * 3 + 2] = best3;
-
-            idx_batch[j * 3 + 0] = besti1;
-            idx_batch[j * 3 + 1] = besti2;
-            idx_batch[j * 3 + 2] = besti3;
-        }
-    }
-}
-
-
-void three_interpolate_cpu(
-    int b, int c, int m, int n,
-    const float *points, 
-    const int *idx,
-    const float *weight, 
-    float *out
-) {
-    for (int batch_index = 0; batch_index < b; ++batch_index) {
-        const float *points_batch = points + batch_index * m * c;
-        const int *idx_batch = idx + batch_index * n * 3;
-        const float *weight_batch = weight + batch_index * n * 3;
-        float *out_batch = out + batch_index * n * c;
-
-        for (int i = 0; i < c * n; ++i) {
-            const int l = i / n;
-            const int j = i % n;
-
-            float w1 = weight_batch[j * 3 + 0];
-            float w2 = weight_batch[j * 3 + 1];
-            float w3 = weight_batch[j * 3 + 2];
-
-            int i1 = idx_batch[j * 3 + 0];
-            int i2 = idx_batch[j * 3 + 1];
-            int i3 = idx_batch[j * 3 + 2];
-
-            out_batch[i] = w1 * points_batch[l * m + i1] +
-                           w2 * points_batch[l * m + i2] +
-                           w3 * points_batch[l * m + i3];
-        }
-    }
-}
-
-
-void three_interpolate_grad_cpu(
-    int b, int c, int n, int m,
-    const float *grad_out,
-    const int *idx, const float *weight,
-    float *grad_points
-) {
-    for (int batch_index = 0; batch_index < b; ++batch_index) {
-        const float *grad_out_batch = grad_out + batch_index * n * c;
-        const int *idx_batch = idx + batch_index * n * 3;
-        const float *weight_batch = weight + batch_index * n * 3;
-        float *grad_points_batch = grad_points + batch_index * m * c;
-
-        // 累积初始值置0
-        for (int i = 0; i < m * c; ++i) {
-            grad_points_batch[i] = 0.0f;
-        }
-
-        for (int i = 0; i < c * n; ++i) {
-            const int l = i / n;
-            const int j = i % n;
-
-            float w1 = weight_batch[j * 3 + 0];
-            float w2 = weight_batch[j * 3 + 1];
-            float w3 = weight_batch[j * 3 + 2];
-
-            int i1 = idx_batch[j * 3 + 0];
-            int i2 = idx_batch[j * 3 + 1];
-            int i3 = idx_batch[j * 3 + 2];
-
-            grad_points_batch[l * m + i1] += w1 * grad_out_batch[i];
-            grad_points_batch[l * m + i2] += w2 * grad_out_batch[i];
-            grad_points_batch[l * m + i3] += w3 * grad_out_batch[i];
-        }
-    }
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/interpolate_gpu.cu

@@ -1,159 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#include <math.h>
-#include <stdio.h>
-#include <stdlib.h>
-
-#include "cuda_utils.h"
-
-// input: unknown(b, n, 3) known(b, m, 3)
-// output: dist2(b, n, 3), idx(b, n, 3)
-__global__ void three_nn_kernel(int b, int n, int m,
-                                const float *__restrict__ unknown,
-                                const float *__restrict__ known,
-                                float *__restrict__ dist2,
-                                int *__restrict__ idx) {
-  int batch_index = blockIdx.x;
-  unknown += batch_index * n * 3;
-  known += batch_index * m * 3;
-  dist2 += batch_index * n * 3;
-  idx += batch_index * n * 3;
-
-  int index = threadIdx.x;
-  int stride = blockDim.x;
-  for (int j = index; j < n; j += stride) {
-    float ux = unknown[j * 3 + 0];
-    float uy = unknown[j * 3 + 1];
-    float uz = unknown[j * 3 + 2];
-
-    double best1 = 1e40, best2 = 1e40, best3 = 1e40;
-    int besti1 = 0, besti2 = 0, besti3 = 0;
-    for (int k = 0; k < m; ++k) {
-      float x = known[k * 3 + 0];
-      float y = known[k * 3 + 1];
-      float z = known[k * 3 + 2];
-      float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z);
-      if (d < best1) {
-        best3 = best2;
-        besti3 = besti2;
-        best2 = best1;
-        besti2 = besti1;
-        best1 = d;
-        besti1 = k;
-      } else if (d < best2) {
-        best3 = best2;
-        besti3 = besti2;
-        best2 = d;
-        besti2 = k;
-      } else if (d < best3) {
-        best3 = d;
-        besti3 = k;
-      }
-    }
-    dist2[j * 3 + 0] = best1;
-    dist2[j * 3 + 1] = best2;
-    dist2[j * 3 + 2] = best3;
-
-    idx[j * 3 + 0] = besti1;
-    idx[j * 3 + 1] = besti2;
-    idx[j * 3 + 2] = besti3;
-  }
-}
-
-void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown,
-                             const float *known, float *dist2, int *idx) {
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  three_nn_kernel<<<b, opt_n_threads(n), 0, stream>>>(b, n, m, unknown, known,
-                                                      dist2, idx);
-
-  CUDA_CHECK_ERRORS();
-}
-
-// input: points(b, c, m), idx(b, n, 3), weight(b, n, 3)
-// output: out(b, c, n)
-__global__ void three_interpolate_kernel(int b, int c, int m, int n,
-                                         const float *__restrict__ points,
-                                         const int *__restrict__ idx,
-                                         const float *__restrict__ weight,
-                                         float *__restrict__ out) {
-  int batch_index = blockIdx.x;
-  points += batch_index * m * c;
-
-  idx += batch_index * n * 3;
-  weight += batch_index * n * 3;
-
-  out += batch_index * n * c;
-
-  const int index = threadIdx.y * blockDim.x + threadIdx.x;
-  const int stride = blockDim.y * blockDim.x;
-  for (int i = index; i < c * n; i += stride) {
-    const int l = i / n;
-    const int j = i % n;
-    float w1 = weight[j * 3 + 0];
-    float w2 = weight[j * 3 + 1];
-    float w3 = weight[j * 3 + 2];
-
-    int i1 = idx[j * 3 + 0];
-    int i2 = idx[j * 3 + 1];
-    int i3 = idx[j * 3 + 2];
-
-    out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 +
-             points[l * m + i3] * w3;
-  }
-}
-
-void three_interpolate_kernel_wrapper(int b, int c, int m, int n,
-                                      const float *points, const int *idx,
-                                      const float *weight, float *out) {
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  three_interpolate_kernel<<<b, opt_block_config(n, c), 0, stream>>>(
-      b, c, m, n, points, idx, weight, out);
-
-  CUDA_CHECK_ERRORS();
-}
-
-// input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3)
-// output: grad_points(b, c, m)
-
-__global__ void three_interpolate_grad_kernel(
-    int b, int c, int n, int m, const float *__restrict__ grad_out,
-    const int *__restrict__ idx, const float *__restrict__ weight,
-    float *__restrict__ grad_points) {
-  int batch_index = blockIdx.x;
-  grad_out += batch_index * n * c;
-  idx += batch_index * n * 3;
-  weight += batch_index * n * 3;
-  grad_points += batch_index * m * c;
-
-  const int index = threadIdx.y * blockDim.x + threadIdx.x;
-  const int stride = blockDim.y * blockDim.x;
-  for (int i = index; i < c * n; i += stride) {
-    const int l = i / n;
-    const int j = i % n;
-    float w1 = weight[j * 3 + 0];
-    float w2 = weight[j * 3 + 1];
-    float w3 = weight[j * 3 + 2];
-
-    int i1 = idx[j * 3 + 0];
-    int i2 = idx[j * 3 + 1];
-    int i3 = idx[j * 3 + 2];
-
-    atomicAdd(grad_points + l * m + i1, grad_out[i] * w1);
-    atomicAdd(grad_points + l * m + i2, grad_out[i] * w2);
-    atomicAdd(grad_points + l * m + i3, grad_out[i] * w3);
-  }
-}
-
-void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m,
-                                           const float *grad_out,
-                                           const int *idx, const float *weight,
-                                           float *grad_points) {
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  three_interpolate_grad_kernel<<<b, opt_block_config(n, c), 0, stream>>>(
-      b, c, n, m, grad_out, idx, weight, grad_points);
-
-  CUDA_CHECK_ERRORS();
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/sampling.cpp

@@ -1,207 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#include "sampling.h"
-#include "utils.h"
-
-
-#ifdef WITH_CUDA
-void gather_points_kernel_wrapper(int b, int c, int n, int npoints,
-                                  const float *points, const int *idx,
-                                  float *out);
-void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
-                                       const float *grad_out, const int *idx,
-                                       float *grad_points);
-
-void furthest_point_sampling_kernel_wrapper(int b, int n, int m,
-                                            const float *dataset, float *temp,
-                                            int *idxs);
-#endif
-
-
-void gather_points_cpu(int b, int c, int n, int npoints, 
-                        const float *points, const int *idx,
-                        float *out);
-void gather_points_grad_cpu(int b, int c, int n, int npoints,
-                            const float *grad_out, const int *idx,
-                            float *grad_points);
-
-void furthest_point_sampling_cpu(int b, int n, int m,
-                                  const float *dataset, float *temp,
-                                  int *idxs);
-
-
-#ifdef DEBUG
-void gather_points_debug(at::Tensor points, at::Tensor idx) {
-    std::cout << "===== sampling.cpp 调试信息 =====" << std::endl;
-    std::cout << "=========== gather_points ===========" << std::endl;
-
-    std::cout << "points.size(): (";
-    for (auto s : points.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" <<std::endl;
-    
-    std::cout << "idx.size():    (";
-    for (auto s : idx.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-
-    std::cout << "====================================" << std::endl;
-}
-
-void gather_points_grad_debug(at::Tensor grad_out, at::Tensor idx,
-                              const int n) {
-    std::cout << "===== sampling.cpp 调试信息 =====" << std::endl;
-    std::cout << "=========== gather_points_grad ===========" << std::endl;
-
-    std::cout << "grad_out.size(): (";
-    for (auto s : grad_out.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" <<std::endl;
-    
-    std::cout << "idx.size():    (";
-    for (auto s : idx.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" << std::endl;
-    
-    std::cout << "n:             " << n << std::endl;
-    std::cout << "====================================" << std::endl;
-}
-
-void furthest_point_sampling_debug(at::Tensor points, const int nsamples) {
-    std::cout << "===== sampling.cpp 调试信息 =====" << std::endl;
-    std::cout << "=========== furthest_point_sampling ===========" << std::endl;
-
-    std::cout << "points.size(): (";
-    for (auto s : points.sizes()) {
-        std::cout << s << " ";
-    }
-    std::cout << ")" <<std::endl;
-    
-    std::cout << "nsamples:      " << nsamples << std::endl;
-    std::cout << "====================================" << std::endl;
-}
-#endif
-
-
-at::Tensor gather_points(at::Tensor points, at::Tensor idx) {
-  CHECK_CONTIGUOUS(points);
-  CHECK_CONTIGUOUS(idx);
-  CHECK_IS_FLOAT(points);
-  CHECK_IS_INT(idx);
-
-  #ifdef WITH_CUDA
-  if (points.is_cuda()) {
-    CHECK_CUDA(idx);
-  }
-  #endif
-
-  #ifdef DEBUG
-  gather_points_debug(points, idx);
-  #endif
-
-  at::Tensor output =
-      torch::zeros({points.size(0), points.size(1), idx.size(1)},
-                   at::device(points.device()).dtype(at::ScalarType::Float));
-
-  #ifdef WITH_CUDA
-  if (points.is_cuda()) {
-    gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2),
-                                 idx.size(1), points.data_ptr<float>(),
-                                 idx.data_ptr<int>(), output.data_ptr<float>());
-  } else {
-  #endif
-  
-    // TORCH_CHECK(false, "CPU not supported");
-    gather_points_cpu(points.size(0), points.size(1), points.size(2),
-                      idx.size(1), points.data_ptr<float>(),
-                      idx.data_ptr<int>(), output.data_ptr<float>());
-  #ifdef WITH_CUDA
-  }
-  #endif
-
-  return output;
-}
-
-at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx,
-                              const int n) {
-  CHECK_CONTIGUOUS(grad_out);
-  CHECK_CONTIGUOUS(idx);
-  CHECK_IS_FLOAT(grad_out);
-  CHECK_IS_INT(idx);
-
-  #ifdef WITH_CUDA
-  if (grad_out.is_cuda()) {
-    CHECK_CUDA(idx);
-  }
-  #endif
-
-  #ifdef DEBUG
-  gather_points_grad_debug(grad_out, idx, n);
-  #endif
-
-  at::Tensor output =
-      torch::zeros({grad_out.size(0), grad_out.size(1), n},
-                   at::device(grad_out.device()).dtype(at::ScalarType::Float));
-
-  #ifdef WITH_CUDA
-  if (grad_out.is_cuda()) {
-    gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n,
-                                      idx.size(1), grad_out.data_ptr<float>(),
-                                      idx.data_ptr<int>(), output.data_ptr<float>());
-  } else {
-  #endif
-
-    // TORCH_CHECK(false, "CPU not supported");
-    gather_points_grad_cpu(grad_out.size(0), grad_out.size(1), n,
-                           idx.size(1), grad_out.data_ptr<float>(),
-                           idx.data_ptr<int>(), output.data_ptr<float>());
-
-  #ifdef WITH_CUDA
-  }
-  #endif
-
-  return output;
-}
-
-at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) {
-  CHECK_CONTIGUOUS(points);
-  CHECK_IS_FLOAT(points);
-
-  #ifdef DEBUG
-  furthest_point_sampling_debug(points, nsamples);
-  #endif
-
-  at::Tensor output =
-      torch::zeros({points.size(0), nsamples},
-                   at::device(points.device()).dtype(at::ScalarType::Int));
-
-  at::Tensor tmp =
-      torch::full({points.size(0), points.size(1)}, 1e10,
-                  at::device(points.device()).dtype(at::ScalarType::Float));
-
-  #ifdef WITH_CUDA
-  if (points.is_cuda()) {
-    furthest_point_sampling_kernel_wrapper(
-        points.size(0), points.size(1), nsamples, points.data_ptr<float>(),
-        tmp.data_ptr<float>(), output.data_ptr<int>());
-  } else {
-  #endif
-
-    // TORCH_CHECK(false, "CPU not supported");
-    furthest_point_sampling_cpu(
-        points.size(0), points.size(1), nsamples, points.data_ptr<float>(), 
-        tmp.data_ptr<float>(), output.data_ptr<int>());
-        
-  #ifdef WITH_CUDA
-  }
-  #endif
-  
-  return output;
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/sampling_cpu.cpp

@@ -1,98 +0,0 @@
-void gather_points_cpu(
-    int b, int c, int n, int npoints, 
-    const float *points, 
-    const int *idx,
-    float *out
-) {
-    for (int i = 0; i < b; ++i) {
-        for (int l = 0; l < c; ++l) {
-            for (int j = 0; j < npoints; ++j) {
-                int a = idx[i * npoints + j];
-                out[(i * c + l) * npoints + j] = points[(i * c + l) * n + a];
-            }
-        }
-    }
-}
-void gather_points_grad_cpu(
-    int b, int c, int n, int npoints,
-    const float *grad_out, 
-    const int *idx,
-    float *grad_points
-) {
-    for (int i = 0; i < b; ++i) {
-        for (int l = 0; l < c; ++l) {
-            for (int j = 0; j < npoints; ++j) {
-                int a = idx[i * npoints + j];
-                grad_points[(i * c + l) * n + a] += grad_out[(i * c + l) * npoints + j];
-            }
-        }
-    }
-}
-
-void furthest_point_sampling_cpu(
-    int b, int n, int m,
-    const float *dataset, 
-    float *temp,
-    int *idxs
-) {
-    // 1. 遍历每个批次（对应CUDA的blockIdx.x）
-    for (int batch_index = 0; batch_index < b; ++batch_index) {
-        const float *dataset_batch = dataset + batch_index * n * 3;
-        float *temp_batch = temp + batch_index * n;
-        int *idxs_batch = idxs + batch_index * m;
-
-        // 2. 初始化：第一个采样点为索引0（对应CUDA的threadIdx.x==0时的赋值）
-        int old = 0;
-        idxs_batch[0] = old;
-
-        // 3. 初始化cur_temp：设为极大值（表示初始时与已选点集（仅一个点）的距离未计算）
-        for (int k = 0; k < n; ++k) {
-            temp_batch[k] = 1e10f; // 极大值，保证首次min操作有效
-        }
-
-        // 4. 迭代m-1次，选择剩余m-1个采样点（对应CUDA的for (int j=1; j<m; j++)）
-        for (int j = 1; j < m; ++j) {
-            // 4.1 提取上一个选中点的坐标（x1,y1,z1）
-            float x1 = dataset_batch[old * 3 + 0];
-            float y1 = dataset_batch[old * 3 + 1];
-            float z1 = dataset_batch[old * 3 + 2];
-
-            // 4.2 遍历所有点，计算距离并更新cur_temp（替换CUDA的线程步长循环）
-            for (int k = 0; k < n; ++k) {
-                // 提取当前点的坐标（x2,y2,z2）
-                float x2 = dataset_batch[k * 3 + 0];
-                float y2 = dataset_batch[k * 3 + 1];
-                float z2 = dataset_batch[k * 3 + 2];
-
-                // 过滤无效点（模长过小，近似为原点，对应CUDA的mag<=1e-3判断）
-                float mag = x2 * x2 + y2 * y2 + z2 * z2;
-                if (mag <= 1e-3f) continue;
-
-                // 计算当前点到上一个选中点的欧氏距离平方（避免开方，不影响大小比较）
-                float dx = x2 - x1;
-                float dy = y2 - y1;
-                float dz = z2 - z1;
-                float d = dx * dx + dy * dy + dz * dz;
-
-                // 核心：更新当前点到已选点集的最小距离
-                if (d < temp_batch[k]) {
-                    temp_batch[k] = d;
-                }
-            }
-
-            // 4.3 规约操作：找到cur_temp中最大值对应的点索引（替换CUDA的__update+共享内存规约）
-            float best = -1.0f;
-            int besti = 0;
-            for (int k = 0; k < n; ++k) {
-                if (temp_batch[k] > best) {
-                    best = temp_batch[k];
-                    besti = k;
-                }
-            }
-
-            // 4.4 记录当前选中点，更新old为下一轮的上一个选中点
-            old = besti;
-            idxs_batch[j] = old;
-        }
-    }
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/_ext_src/src/sampling_gpu.cu

@@ -1,234 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-// 
-// This source code is licensed under the MIT license found in the
-// LICENSE file in the root directory of this source tree.
-
-#include <stdio.h>
-#include <stdlib.h>
-
-#include "cuda_utils.h"
-
-// input: points(b, c, n) idx(b, m)
-// output: out(b, c, m)
-__global__ void gather_points_kernel(int b, int c, int n, int m,
-                                     const float *__restrict__ points,
-                                     const int *__restrict__ idx,
-                                     float *__restrict__ out) {
-  for (int i = blockIdx.x; i < b; i += gridDim.x) {
-    for (int l = blockIdx.y; l < c; l += gridDim.y) {
-      for (int j = threadIdx.x; j < m; j += blockDim.x) {
-        int a = idx[i * m + j];
-        out[(i * c + l) * m + j] = points[(i * c + l) * n + a];
-      }
-    }
-  }
-}
-
-void gather_points_kernel_wrapper(int b, int c, int n, int npoints,
-                                  const float *points, const int *idx,
-                                  float *out) {
-  gather_points_kernel<<<dim3(b, c, 1), opt_n_threads(npoints), 0,
-                         at::cuda::getCurrentCUDAStream()>>>(b, c, n, npoints,
-                                                             points, idx, out);
-
-  CUDA_CHECK_ERRORS();
-}
-
-// input: grad_out(b, c, m) idx(b, m)
-// output: grad_points(b, c, n)
-__global__ void gather_points_grad_kernel(int b, int c, int n, int m,
-                                          const float *__restrict__ grad_out,
-                                          const int *__restrict__ idx,
-                                          float *__restrict__ grad_points) {
-  for (int i = blockIdx.x; i < b; i += gridDim.x) {
-    for (int l = blockIdx.y; l < c; l += gridDim.y) {
-      for (int j = threadIdx.x; j < m; j += blockDim.x) {
-        int a = idx[i * m + j];
-        atomicAdd(grad_points + (i * c + l) * n + a,
-                  grad_out[(i * c + l) * m + j]);
-      }
-    }
-  }
-}
-
-void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
-                                       const float *grad_out, const int *idx,
-                                       float *grad_points) {
-  gather_points_grad_kernel<<<dim3(b, c, 1), opt_n_threads(npoints), 0,
-                              at::cuda::getCurrentCUDAStream()>>>(
-      b, c, n, npoints, grad_out, idx, grad_points);
-
-  CUDA_CHECK_ERRORS();
-}
-
-__device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i,
-                         int idx1, int idx2) {
-  const float v1 = dists[idx1], v2 = dists[idx2];
-  const int i1 = dists_i[idx1], i2 = dists_i[idx2];
-  dists[idx1] = max(v1, v2);
-  dists_i[idx1] = v2 > v1 ? i2 : i1;
-}
-
-// Input dataset: (b, n, 3), tmp: (b, n)
-// Ouput idxs (b, m)
-template <unsigned int block_size>
-__global__ void furthest_point_sampling_kernel(
-    int b, int n, int m, const float *__restrict__ dataset,
-    float *__restrict__ temp, int *__restrict__ idxs) {
-  if (m <= 0) return;
-  __shared__ float dists[block_size];
-  __shared__ int dists_i[block_size];
-
-  int batch_index = blockIdx.x;
-  dataset += batch_index * n * 3;
-  temp += batch_index * n;
-  idxs += batch_index * m;
-
-  int tid = threadIdx.x;
-  const int stride = block_size;
-
-  int old = 0;
-  if (threadIdx.x == 0) idxs[0] = old;
-
-  __syncthreads();
-  for (int j = 1; j < m; j++) {
-    int besti = 0;
-    float best = -1;
-    float x1 = dataset[old * 3 + 0];
-    float y1 = dataset[old * 3 + 1];
-    float z1 = dataset[old * 3 + 2];
-    for (int k = tid; k < n; k += stride) {
-      float x2, y2, z2;
-      x2 = dataset[k * 3 + 0];
-      y2 = dataset[k * 3 + 1];
-      z2 = dataset[k * 3 + 2];
-      float mag = (x2 * x2) + (y2 * y2) + (z2 * z2);
-      if (mag <= 1e-3) continue;
-
-      float d =
-          (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1);
-
-      float d2 = min(d, temp[k]);
-      temp[k] = d2;
-      besti = d2 > best ? k : besti;
-      best = d2 > best ? d2 : best;
-    }
-    dists[tid] = best;
-    dists_i[tid] = besti;
-    __syncthreads();
-
-    if (block_size >= 512) {
-      if (tid < 256) {
-        __update(dists, dists_i, tid, tid + 256);
-      }
-      __syncthreads();
-    }
-    if (block_size >= 256) {
-      if (tid < 128) {
-        __update(dists, dists_i, tid, tid + 128);
-      }
-      __syncthreads();
-    }
-    if (block_size >= 128) {
-      if (tid < 64) {
-        __update(dists, dists_i, tid, tid + 64);
-      }
-      __syncthreads();
-    }
-    if (block_size >= 64) {
-      if (tid < 32) {
-        __update(dists, dists_i, tid, tid + 32);
-      }
-      __syncthreads();
-    }
-    if (block_size >= 32) {
-      if (tid < 16) {
-        __update(dists, dists_i, tid, tid + 16);
-      }
-      __syncthreads();
-    }
-    if (block_size >= 16) {
-      if (tid < 8) {
-        __update(dists, dists_i, tid, tid + 8);
-      }
-      __syncthreads();
-    }
-    if (block_size >= 8) {
-      if (tid < 4) {
-        __update(dists, dists_i, tid, tid + 4);
-      }
-      __syncthreads();
-    }
-    if (block_size >= 4) {
-      if (tid < 2) {
-        __update(dists, dists_i, tid, tid + 2);
-      }
-      __syncthreads();
-    }
-    if (block_size >= 2) {
-      if (tid < 1) {
-        __update(dists, dists_i, tid, tid + 1);
-      }
-      __syncthreads();
-    }
-
-    old = dists_i[0];
-    if (tid == 0) idxs[j] = old;
-  }
-}
-
-void furthest_point_sampling_kernel_wrapper(int b, int n, int m,
-                                            const float *dataset, float *temp,
-                                            int *idxs) {
-  unsigned int n_threads = opt_n_threads(n);
-
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  switch (n_threads) {
-    case 512:
-      furthest_point_sampling_kernel<512>
-          <<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
-      break;
-    case 256:
-      furthest_point_sampling_kernel<256>
-          <<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
-      break;
-    case 128:
-      furthest_point_sampling_kernel<128>
-          <<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
-      break;
-    case 64:
-      furthest_point_sampling_kernel<64>
-          <<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
-      break;
-    case 32:
-      furthest_point_sampling_kernel<32>
-          <<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
-      break;
-    case 16:
-      furthest_point_sampling_kernel<16>
-          <<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
-      break;
-    case 8:
-      furthest_point_sampling_kernel<8>
-          <<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
-      break;
-    case 4:
-      furthest_point_sampling_kernel<4>
-          <<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
-      break;
-    case 2:
-      furthest_point_sampling_kernel<2>
-          <<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
-      break;
-    case 1:
-      furthest_point_sampling_kernel<1>
-          <<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
-      break;
-    default:
-      furthest_point_sampling_kernel<512>
-          <<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
-  }
-
-  CUDA_CHECK_ERRORS();
-}

vision_detect/vision_detect/VisionDetect/net/pointnet2/pointnet2/__init__.py

@@ -1 +0,0 @@
-from .pointnet2 import *

vision_detect/vision_detect/VisionDetect/net/pointnet2/pointnet2_modules.py

@@ -1,518 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-# 
-# This source code is licensed under the MIT license found in the
-# LICENSE file in the root directory of this source tree.
-
-''' Pointnet2 layers.
-Modified based on: https://github.com/erikwijmans/Pointnet2_PyTorch
-Extended with the following:
-1. Uniform sampling in each local region (sample_uniformly)
-2. Return sampled points indices to support votenet.
-'''
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-# import os
-# import sys
-# BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-# sys.path.append(BASE_DIR)
-
-from . import pointnet2_utils
-from . import pytorch_utils as pt_utils
-from typing import List
-
-
-class _PointnetSAModuleBase(nn.Module):
-
-    def __init__(self):
-        super().__init__()
-        self.npoint = None
-        self.groupers = None
-        self.mlps = None
-
-    def forward(self, xyz: torch.Tensor,
-                features: torch.Tensor = None) -> (torch.Tensor, torch.Tensor):
-        r"""
-        Parameters
-        ----------
-        xyz : torch.Tensor
-            (B, N, 3) tensor of the xyz coordinates of the features
-        features : torch.Tensor
-            (B, N, C) tensor of the descriptors of the the features
-
-        Returns
-        -------
-        new_xyz : torch.Tensor
-            (B, npoint, 3) tensor of the new features' xyz
-        new_features : torch.Tensor
-            (B, npoint, \sum_k(mlps[k][-1])) tensor of the new_features descriptors
-        """
-
-        new_features_list = []
-
-        xyz_flipped = xyz.transpose(1, 2).contiguous()
-        new_xyz = pointnet2_utils.gather_operation(
-            xyz_flipped,
-            pointnet2_utils.furthest_point_sample(xyz, self.npoint)
-        ).transpose(1, 2).contiguous() if self.npoint is not None else None
-
-        for i in range(len(self.groupers)):
-            new_features = self.groupers[i](
-                xyz, new_xyz, features
-            )  # (B, C, npoint, nsample)
-
-            new_features = self.mlps[i](
-                new_features
-            )  # (B, mlp[-1], npoint, nsample)
-            new_features = F.max_pool2d(
-                new_features, kernel_size=[1, new_features.size(3)]
-            )  # (B, mlp[-1], npoint, 1)
-            new_features = new_features.squeeze(-1)  # (B, mlp[-1], npoint)
-
-            new_features_list.append(new_features)
-
-        return new_xyz, torch.cat(new_features_list, dim=1)
-
-
-class PointnetSAModuleMSG(_PointnetSAModuleBase):
-    r"""Pointnet set abstrction layer with multiscale grouping
-
-    Parameters
-    ----------
-    npoint : int
-        Number of features
-    radii : list of float32
-        list of radii to group with
-    nsamples : list of int32
-        Number of samples in each ball query
-    mlps : list of list of int32
-        Spec of the pointnet before the global max_pool for each scale
-    bn : bool
-        Use batchnorm
-    """
-
-    def __init__(
-            self,
-            *,
-            npoint: int,
-            radii: List[float],
-            nsamples: List[int],
-            mlps: List[List[int]],
-            bn: bool = True,
-            use_xyz: bool = True, 
-            sample_uniformly: bool = False
-    ):
-        super().__init__()
-
-        assert len(radii) == len(nsamples) == len(mlps)
-
-        self.npoint = npoint
-        self.groupers = nn.ModuleList()
-        self.mlps = nn.ModuleList()
-        for i in range(len(radii)):
-            radius = radii[i]
-            nsample = nsamples[i]
-            self.groupers.append(
-                pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz, sample_uniformly=sample_uniformly)
-                if npoint is not None else pointnet2_utils.GroupAll(use_xyz)
-            )
-            mlp_spec = mlps[i]
-            if use_xyz:
-                mlp_spec[0] += 3
-
-            self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn))
-
-
-class PointnetSAModule(PointnetSAModuleMSG):
-    r"""Pointnet set abstrction layer
-
-    Parameters
-    ----------
-    npoint : int
-        Number of features
-    radius : float
-        Radius of ball
-    nsample : int
-        Number of samples in the ball query
-    mlp : list
-        Spec of the pointnet before the global max_pool
-    bn : bool
-        Use batchnorm
-    """
-
-    def __init__(
-            self,
-            *,
-            mlp: List[int],
-            npoint: int = None,
-            radius: float = None,
-            nsample: int = None,
-            bn: bool = True,
-            use_xyz: bool = True
-    ):
-        super().__init__(
-            mlps=[mlp],
-            npoint=npoint,
-            radii=[radius],
-            nsamples=[nsample],
-            bn=bn,
-            use_xyz=use_xyz
-        )
-
-
-class PointnetSAModuleVotes(nn.Module):
-    ''' Modified based on _PointnetSAModuleBase and PointnetSAModuleMSG
-    with extra support for returning point indices for getting their GT votes '''
-
-    def __init__(
-            self,
-            *,
-            mlp: List[int],
-            npoint: int = None,
-            radius: float = None,
-            nsample: int = None,
-            bn: bool = True,
-            use_xyz: bool = True,
-            pooling: str = 'max',
-            sigma: float = None, # for RBF pooling
-            normalize_xyz: bool = False, # noramlize local XYZ with radius
-            sample_uniformly: bool = False,
-            ret_unique_cnt: bool = False
-    ):
-        super().__init__()
-
-        self.npoint = npoint
-        self.radius = radius
-        self.nsample = nsample
-        self.pooling = pooling
-        self.mlp_module = None
-        self.use_xyz = use_xyz
-        self.sigma = sigma
-        if self.sigma is None:
-            self.sigma = self.radius/2
-        self.normalize_xyz = normalize_xyz
-        self.ret_unique_cnt = ret_unique_cnt
-
-        if npoint is not None:
-            self.grouper = pointnet2_utils.QueryAndGroup(radius, nsample,
-                use_xyz=use_xyz, ret_grouped_xyz=True, normalize_xyz=normalize_xyz,
-                sample_uniformly=sample_uniformly, ret_unique_cnt=ret_unique_cnt)
-        else:
-            self.grouper = pointnet2_utils.GroupAll(use_xyz, ret_grouped_xyz=True)
-
-        mlp_spec = mlp
-        if use_xyz and len(mlp_spec)>0:
-            mlp_spec[0] += 3
-        self.mlp_module = pt_utils.SharedMLP(mlp_spec, bn=bn)
-
-
-    def forward(self, xyz: torch.Tensor,
-                features: torch.Tensor = None,
-                inds: torch.Tensor = None) -> (torch.Tensor, torch.Tensor):
-        r"""
-        Parameters
-        ----------
-        xyz : torch.Tensor
-            (B, N, 3) tensor of the xyz coordinates of the features
-        features : torch.Tensor
-            (B, C, N) tensor of the descriptors of the the features
-        inds : torch.Tensor
-            (B, npoint) tensor that stores index to the xyz points (values in 0-N-1)
-
-        Returns
-        -------
-        new_xyz : torch.Tensor
-            (B, npoint, 3) tensor of the new features' xyz
-        new_features : torch.Tensor
-            (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors
-        inds: torch.Tensor
-            (B, npoint) tensor of the inds
-        """
-
-        xyz_flipped = xyz.transpose(1, 2).contiguous()
-        if inds is None:
-            inds = pointnet2_utils.furthest_point_sample(xyz, self.npoint)
-        else:
-            assert(inds.shape[1] == self.npoint)
-        new_xyz = pointnet2_utils.gather_operation(
-            xyz_flipped, inds
-        ).transpose(1, 2).contiguous() if self.npoint is not None else None
-
-        if not self.ret_unique_cnt:
-            grouped_features, grouped_xyz = self.grouper(
-                xyz, new_xyz, features
-            )  # (B, C, npoint, nsample)
-        else:
-            grouped_features, grouped_xyz, unique_cnt = self.grouper(
-                xyz, new_xyz, features
-            )  # (B, C, npoint, nsample), (B,3,npoint,nsample), (B,npoint)
-
-        new_features = self.mlp_module(
-            grouped_features
-        )  # (B, mlp[-1], npoint, nsample)
-        if self.pooling == 'max':
-            new_features = F.max_pool2d(
-                new_features, kernel_size=[1, new_features.size(3)]
-            )  # (B, mlp[-1], npoint, 1)
-        elif self.pooling == 'avg':
-            new_features = F.avg_pool2d(
-                new_features, kernel_size=[1, new_features.size(3)]
-            )  # (B, mlp[-1], npoint, 1)
-        elif self.pooling == 'rbf': 
-            # Use radial basis function kernel for weighted sum of features (normalized by nsample and sigma)
-            # Ref: https://en.wikipedia.org/wiki/Radial_basis_function_kernel
-            rbf = torch.exp(-1 * grouped_xyz.pow(2).sum(1,keepdim=False) / (self.sigma**2) / 2) # (B, npoint, nsample)
-            new_features = torch.sum(new_features * rbf.unsqueeze(1), -1, keepdim=True) / float(self.nsample) # (B, mlp[-1], npoint, 1)
-        new_features = new_features.squeeze(-1)  # (B, mlp[-1], npoint)
-
-        if not self.ret_unique_cnt:
-            return new_xyz, new_features, inds
-        else:
-            return new_xyz, new_features, inds, unique_cnt
-
-class PointnetSAModuleMSGVotes(nn.Module):
-    ''' Modified based on _PointnetSAModuleBase and PointnetSAModuleMSG
-    with extra support for returning point indices for getting their GT votes '''
-
-    def __init__(
-            self,
-            *,
-            mlps: List[List[int]],
-            npoint: int,
-            radii: List[float],
-            nsamples: List[int],
-            bn: bool = True,
-            use_xyz: bool = True,
-            sample_uniformly: bool = False
-    ):
-        super().__init__()
-
-        assert(len(mlps) == len(nsamples) == len(radii))
-
-        self.npoint = npoint
-        self.groupers = nn.ModuleList()
-        self.mlps = nn.ModuleList()
-        for i in range(len(radii)):
-            radius = radii[i]
-            nsample = nsamples[i]
-            self.groupers.append(
-                pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz, sample_uniformly=sample_uniformly)
-                if npoint is not None else pointnet2_utils.GroupAll(use_xyz)
-            )
-            mlp_spec = mlps[i]
-            if use_xyz:
-                mlp_spec[0] += 3
-
-            self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn))
-
-    def forward(self, xyz: torch.Tensor,
-                features: torch.Tensor = None, inds: torch.Tensor = None) -> (torch.Tensor, torch.Tensor):
-        r"""
-        Parameters
-        ----------
-        xyz : torch.Tensor
-            (B, N, 3) tensor of the xyz coordinates of the features
-        features : torch.Tensor
-            (B, C, C) tensor of the descriptors of the the features
-        inds : torch.Tensor
-            (B, npoint) tensor that stores index to the xyz points (values in 0-N-1)
-
-        Returns
-        -------
-        new_xyz : torch.Tensor
-            (B, npoint, 3) tensor of the new features' xyz
-        new_features : torch.Tensor
-            (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors
-        inds: torch.Tensor
-            (B, npoint) tensor of the inds
-        """
-        new_features_list = []
-
-        xyz_flipped = xyz.transpose(1, 2).contiguous()
-        if inds is None:
-            inds = pointnet2_utils.furthest_point_sample(xyz, self.npoint)
-        new_xyz = pointnet2_utils.gather_operation(
-            xyz_flipped, inds
-        ).transpose(1, 2).contiguous() if self.npoint is not None else None
-
-        for i in range(len(self.groupers)):
-            new_features = self.groupers[i](
-                xyz, new_xyz, features
-            )  # (B, C, npoint, nsample)
-            new_features = self.mlps[i](
-                new_features
-            )  # (B, mlp[-1], npoint, nsample)
-            new_features = F.max_pool2d(
-                new_features, kernel_size=[1, new_features.size(3)]
-            )  # (B, mlp[-1], npoint, 1)
-            new_features = new_features.squeeze(-1)  # (B, mlp[-1], npoint)
-
-            new_features_list.append(new_features)
-
-        return new_xyz, torch.cat(new_features_list, dim=1), inds
-
-
-class PointnetFPModule(nn.Module):
-    r"""Propigates the features of one set to another
-
-    Parameters
-    ----------
-    mlp : list
-        Pointnet module parameters
-    bn : bool
-        Use batchnorm
-    """
-
-    def __init__(self, *, mlp: List[int], bn: bool = True):
-        super().__init__()
-        self.mlp = pt_utils.SharedMLP(mlp, bn=bn)
-
-    def forward(
-            self, unknown: torch.Tensor, known: torch.Tensor,
-            unknow_feats: torch.Tensor, known_feats: torch.Tensor
-    ) -> torch.Tensor:
-        r"""
-        Parameters
-        ----------
-        unknown : torch.Tensor
-            (B, n, 3) tensor of the xyz positions of the unknown features
-        known : torch.Tensor
-            (B, m, 3) tensor of the xyz positions of the known features
-        unknow_feats : torch.Tensor
-            (B, C1, n) tensor of the features to be propigated to
-        known_feats : torch.Tensor
-            (B, C2, m) tensor of features to be propigated
-
-        Returns
-        -------
-        new_features : torch.Tensor
-            (B, mlp[-1], n) tensor of the features of the unknown features
-        """
-
-        if known is not None:
-            dist, idx = pointnet2_utils.three_nn(unknown, known)
-            dist_recip = 1.0 / (dist + 1e-8)
-            norm = torch.sum(dist_recip, dim=2, keepdim=True)
-            weight = dist_recip / norm
-
-            interpolated_feats = pointnet2_utils.three_interpolate(
-                known_feats, idx, weight
-            )
-        else:
-            interpolated_feats = known_feats.expand(
-                *known_feats.size()[0:2], unknown.size(1)
-            )
-
-        if unknow_feats is not None:
-            new_features = torch.cat([interpolated_feats, unknow_feats],
-                                   dim=1)  #(B, C2 + C1, n)
-        else:
-            new_features = interpolated_feats
-
-        new_features = new_features.unsqueeze(-1)
-        new_features = self.mlp(new_features)
-
-        return new_features.squeeze(-1)
-
-class PointnetLFPModuleMSG(nn.Module):
-    ''' Modified based on _PointnetSAModuleBase and PointnetSAModuleMSG
-    learnable feature propagation layer.'''
-
-    def __init__(
-            self,
-            *,
-            mlps: List[List[int]],
-            radii: List[float],
-            nsamples: List[int],
-            post_mlp: List[int],
-            bn: bool = True,
-            use_xyz: bool = True,
-            sample_uniformly: bool = False
-    ):
-        super().__init__()
-
-        assert(len(mlps) == len(nsamples) == len(radii))
-        
-        self.post_mlp = pt_utils.SharedMLP(post_mlp, bn=bn)
-
-        self.groupers = nn.ModuleList()
-        self.mlps = nn.ModuleList()
-        for i in range(len(radii)):
-            radius = radii[i]
-            nsample = nsamples[i]
-            self.groupers.append(
-                pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz,
-                    sample_uniformly=sample_uniformly)
-            )
-            mlp_spec = mlps[i]
-            if use_xyz:
-                mlp_spec[0] += 3
-
-            self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn))
-
-    def forward(self, xyz2: torch.Tensor, xyz1: torch.Tensor,
-                features2: torch.Tensor, features1: torch.Tensor) -> torch.Tensor:
-        r""" Propagate features from xyz1 to xyz2.
-        Parameters
-        ----------
-        xyz2 : torch.Tensor
-            (B, N2, 3) tensor of the xyz coordinates of the features
-        xyz1 : torch.Tensor
-            (B, N1, 3) tensor of the xyz coordinates of the features
-        features2 : torch.Tensor
-            (B, C2, N2) tensor of the descriptors of the the features
-        features1 : torch.Tensor
-            (B, C1, N1) tensor of the descriptors of the the features
-
-        Returns
-        -------
-        new_features1 : torch.Tensor
-            (B, \sum_k(mlps[k][-1]), N1) tensor of the new_features descriptors
-        """
-        new_features_list = []
-
-        for i in range(len(self.groupers)):
-            new_features = self.groupers[i](
-                xyz1, xyz2, features1
-            )  # (B, C1, N2, nsample)
-            new_features = self.mlps[i](
-                new_features
-            )  # (B, mlp[-1], N2, nsample)
-            new_features = F.max_pool2d(
-                new_features, kernel_size=[1, new_features.size(3)]
-            )  # (B, mlp[-1], N2, 1)
-            new_features = new_features.squeeze(-1)  # (B, mlp[-1], N2)
-
-            if features2 is not None:
-                new_features = torch.cat([new_features, features2],
-                                           dim=1)  #(B, mlp[-1] + C2, N2)
-
-            new_features = new_features.unsqueeze(-1)
-            new_features = self.post_mlp(new_features)
-
-            new_features_list.append(new_features)
-
-        return torch.cat(new_features_list, dim=1).squeeze(-1)
-
-
-if __name__ == "__main__":
-    from torch.autograd import Variable
-    torch.manual_seed(1)
-    torch.cuda.manual_seed_all(1)
-    xyz = Variable(torch.randn(2, 9, 3).cuda(), requires_grad=True)
-    xyz_feats = Variable(torch.randn(2, 9, 6).cuda(), requires_grad=True)
-
-    test_module = PointnetSAModuleMSG(
-        npoint=2, radii=[5.0, 10.0], nsamples=[6, 3], mlps=[[9, 3], [9, 6]]
-    )
-    test_module.cuda()
-    print(test_module(xyz, xyz_feats))
-
-    for _ in range(1):
-        _, new_features = test_module(xyz, xyz_feats)
-        new_features.backward(
-            torch.cuda.FloatTensor(*new_features.size()).fill_(1)
-        )
-        print(new_features)
-        print(xyz.grad)

vision_detect/vision_detect/VisionDetect/net/pointnet2/pointnet2_utils.py

@@ -1,554 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-# 
-# This source code is licensed under the MIT license found in the
-# LICENSE file in the root directory of this source tree.
-
-''' Modified based on: https://github.com/erikwijmans/Pointnet2_PyTorch '''
-from __future__ import (
-    division,
-    absolute_import,
-    with_statement,
-    print_function,
-    unicode_literals,
-)
-import torch
-from torch.autograd import Function
-import torch.nn as nn
-# import pytorch_utils as pt_utils
-from . import pytorch_utils as pt_utils
-# import sys
-
-try:
-    import builtins
-except:
-    import __builtin__ as builtins
-
-try:
-    from . import pointnet2
-except ImportError:
-    if not getattr(builtins, "__POINTNET2_SETUP__", False):
-        raise ImportError(
-            "Could not import pointnet2 module.\n"
-            "Please see the setup instructions in the README: "
-            "https://github.com/erikwijmans/Pointnet2_PyTorch/blob/master/README.rst"
-        )
-
-if False:
-    # Workaround for type hints without depending on the `typing` module
-    from typing import *
-
-
-class RandomDropout(nn.Module):
-    def __init__(self, p=0.5, inplace=False):
-        super(RandomDropout, self).__init__()
-        self.p = p
-        self.inplace = inplace
-
-    def forward(self, X):
-        theta = torch.Tensor(1).uniform_(0, self.p)[0]
-        return pt_utils.feature_dropout_no_scaling(X, theta, self.train, self.inplace)
-
-
-class FurthestPointSampling(Function):
-    @staticmethod
-    def forward(ctx, xyz, npoint):
-        # type: (Any, torch.Tensor, int) -> torch.Tensor
-        r"""
-        Uses iterative furthest point sampling to select a set of npoint features that have the largest
-        minimum distance
-
-        Parameters
-        ----------
-        xyz : torch.Tensor
-            (B, N, 3) tensor where N > npoint
-        npoint : int32
-            number of features in the sampled set
-
-        Returns
-        -------
-        torch.Tensor
-            (B, npoint) tensor containing the set
-        """
-        return pointnet2.furthest_point_sampling(xyz, npoint)
-
-    @staticmethod
-    def backward(xyz, a=None):
-        return None, None
-
-
-furthest_point_sample = FurthestPointSampling.apply
-
-
-class GatherOperation(Function):
-    @staticmethod
-    def forward(ctx, features, idx):
-        # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor
-        r"""
-
-        Parameters
-        ----------
-        features : torch.Tensor
-            (B, C, N) tensor
-
-        idx : torch.Tensor
-            (B, npoint) tensor of the features to gather
-
-        Returns
-        -------
-        torch.Tensor
-            (B, C, npoint) tensor
-        """
-
-        _, C, N = features.size()
-
-        ctx.for_backwards = (idx, C, N)
-
-        return pointnet2.gather_points(features, idx)
-
-    @staticmethod
-    def backward(ctx, grad_out):
-        idx, C, N = ctx.for_backwards
-
-        grad_features = pointnet2.gather_points_grad(grad_out.contiguous(), idx, N)
-        return grad_features, None
-
-
-gather_operation = GatherOperation.apply
-
-
-class ThreeNN(Function):
-    @staticmethod
-    def forward(ctx, unknown, known):
-        # type: (Any, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]
-        r"""
-            Find the three nearest neighbors of unknown in known
-        Parameters
-        ----------
-        unknown : torch.Tensor
-            (B, n, 3) tensor of known features
-        known : torch.Tensor
-            (B, m, 3) tensor of unknown features
-
-        Returns
-        -------
-        dist : torch.Tensor
-            (B, n, 3) l2 distance to the three nearest neighbors
-        idx : torch.Tensor
-            (B, n, 3) index of 3 nearest neighbors
-        """
-        dist2, idx = pointnet2.three_nn(unknown, known)
-
-        return torch.sqrt(dist2), idx
-
-    @staticmethod
-    def backward(ctx, a=None, b=None):
-        return None, None
-
-
-three_nn = ThreeNN.apply
-
-
-class ThreeInterpolate(Function):
-    @staticmethod
-    def forward(ctx, features, idx, weight):
-        # type(Any, torch.Tensor, torch.Tensor, torch.Tensor) -> Torch.Tensor
-        r"""
-            Performs weight linear interpolation on 3 features
-        Parameters
-        ----------
-        features : torch.Tensor
-            (B, c, m) Features descriptors to be interpolated from
-        idx : torch.Tensor
-            (B, n, 3) three nearest neighbors of the target features in features
-        weight : torch.Tensor
-            (B, n, 3) weights
-
-        Returns
-        -------
-        torch.Tensor
-            (B, c, n) tensor of the interpolated features
-        """
-        B, c, m = features.size()
-        n = idx.size(1)
-
-        ctx.three_interpolate_for_backward = (idx, weight, m)
-
-        return pointnet2.three_interpolate(features, idx, weight)
-
-    @staticmethod
-    def backward(ctx, grad_out):
-        # type: (Any, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
-        r"""
-        Parameters
-        ----------
-        grad_out : torch.Tensor
-            (B, c, n) tensor with gradients of ouputs
-
-        Returns
-        -------
-        grad_features : torch.Tensor
-            (B, c, m) tensor with gradients of features
-
-        None
-
-        None
-        """
-        idx, weight, m = ctx.three_interpolate_for_backward
-
-        grad_features = pointnet2.three_interpolate_grad(
-            grad_out.contiguous(), idx, weight, m
-        )
-
-        return grad_features, None, None
-
-
-three_interpolate = ThreeInterpolate.apply
-
-
-class GroupingOperation(Function):
-    @staticmethod
-    def forward(ctx, features, idx):
-        # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor
-        r"""
-
-        Parameters
-        ----------
-        features : torch.Tensor
-            (B, C, N) tensor of features to group
-        idx : torch.Tensor
-            (B, npoint, nsample) tensor containing the indicies of features to group with
-
-        Returns
-        -------
-        torch.Tensor
-            (B, C, npoint, nsample) tensor
-        """
-        B, nfeatures, nsample = idx.size()
-        _, C, N = features.size()
-
-        ctx.for_backwards = (idx, N)
-
-        return pointnet2.group_points(features, idx)
-
-    @staticmethod
-    def backward(ctx, grad_out):
-        # type: (Any, torch.tensor) -> Tuple[torch.Tensor, torch.Tensor]
-        r"""
-
-        Parameters
-        ----------
-        grad_out : torch.Tensor
-            (B, C, npoint, nsample) tensor of the gradients of the output from forward
-
-        Returns
-        -------
-        torch.Tensor
-            (B, C, N) gradient of the features
-        None
-        """
-        idx, N = ctx.for_backwards
-
-        grad_features = pointnet2.group_points_grad(grad_out.contiguous(), idx, N)
-
-        return grad_features, None
-
-
-grouping_operation = GroupingOperation.apply
-
-
-class BallQuery(Function):
-    @staticmethod
-    def forward(ctx, radius, nsample, xyz, new_xyz):
-        # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor
-        r"""
-
-        Parameters
-        ----------
-        radius : float
-            radius of the balls
-        nsample : int
-            maximum number of features in the balls
-        xyz : torch.Tensor
-            (B, N, 3) xyz coordinates of the features
-        new_xyz : torch.Tensor
-            (B, npoint, 3) centers of the ball query
-
-        Returns
-        -------
-        torch.Tensor
-            (B, npoint, nsample) tensor with the indicies of the features that form the query balls
-        """
-        return pointnet2.ball_query(new_xyz, xyz, radius, nsample)
-
-    @staticmethod
-    def backward(ctx, a=None):
-        return None, None, None, None
-
-
-ball_query = BallQuery.apply
-
-
-class QueryAndGroup(nn.Module):
-    r"""
-    Groups with a ball query of radius
-
-    Parameters
-    ---------
-    radius : float32
-        Radius of ball
-    nsample : int32
-        Maximum number of features to gather in the ball
-    """
-
-    def __init__(self, radius, nsample, use_xyz=True, ret_grouped_xyz=False, normalize_xyz=False, sample_uniformly=False, ret_unique_cnt=False):
-        # type: (QueryAndGroup, float, int, bool) -> None
-        super(QueryAndGroup, self).__init__()
-        self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz
-        self.ret_grouped_xyz = ret_grouped_xyz
-        self.normalize_xyz = normalize_xyz
-        self.sample_uniformly = sample_uniformly
-        self.ret_unique_cnt = ret_unique_cnt
-        if self.ret_unique_cnt:
-            assert(self.sample_uniformly)
-
-    def forward(self, xyz, new_xyz, features=None):
-        # type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[Torch.Tensor]
-        r"""
-        Parameters
-        ----------
-        xyz : torch.Tensor
-            xyz coordinates of the features (B, N, 3)
-        new_xyz : torch.Tensor
-            centriods (B, npoint, 3)
-        features : torch.Tensor
-            Descriptors of the features (B, C, N)
-
-        Returns
-        -------
-        new_features : torch.Tensor
-            (B, 3 + C, npoint, nsample) tensor
-        """
-        idx = ball_query(self.radius, self.nsample, xyz, new_xyz)
-
-        if self.sample_uniformly:
-            unique_cnt = torch.zeros((idx.shape[0], idx.shape[1]))
-            for i_batch in range(idx.shape[0]):
-                for i_region in range(idx.shape[1]):
-                    unique_ind = torch.unique(idx[i_batch, i_region, :])
-                    num_unique = unique_ind.shape[0]
-                    unique_cnt[i_batch, i_region] = num_unique
-                    sample_ind = torch.randint(0, num_unique, (self.nsample - num_unique,), dtype=torch.long)
-                    all_ind = torch.cat((unique_ind, unique_ind[sample_ind]))
-                    idx[i_batch, i_region, :] = all_ind
-
-
-        xyz_trans = xyz.transpose(1, 2).contiguous()
-        grouped_xyz = grouping_operation(xyz_trans, idx)  # (B, 3, npoint, nsample)
-        grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1)
-        if self.normalize_xyz:
-            grouped_xyz /= self.radius
-
-        if features is not None:
-            grouped_features = grouping_operation(features, idx)
-            if self.use_xyz:
-                new_features = torch.cat(
-                    [grouped_xyz, grouped_features], dim=1
-                )  # (B, C + 3, npoint, nsample)
-            else:
-                new_features = grouped_features
-        else:
-            assert (
-                self.use_xyz
-            ), "Cannot have not features and not use xyz as a feature!"
-            new_features = grouped_xyz
-
-        ret = [new_features]
-        if self.ret_grouped_xyz:
-            ret.append(grouped_xyz)
-        if self.ret_unique_cnt:
-            ret.append(unique_cnt)
-        if len(ret) == 1:
-            return ret[0]
-        else:
-            return tuple(ret)
-
-
-class GroupAll(nn.Module):
-    r"""
-    Groups all features
-
-    Parameters
-    ---------
-    """
-
-    def __init__(self, use_xyz=True, ret_grouped_xyz=False):
-        # type: (GroupAll, bool) -> None
-        super(GroupAll, self).__init__()
-        self.use_xyz = use_xyz
-
-    def forward(self, xyz, new_xyz, features=None):
-        # type: (GroupAll, torch.Tensor, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor]
-        r"""
-        Parameters
-        ----------
-        xyz : torch.Tensor
-            xyz coordinates of the features (B, N, 3)
-        new_xyz : torch.Tensor
-            Ignored
-        features : torch.Tensor
-            Descriptors of the features (B, C, N)
-
-        Returns
-        -------
-        new_features : torch.Tensor
-            (B, C + 3, 1, N) tensor
-        """
-
-        grouped_xyz = xyz.transpose(1, 2).unsqueeze(2)
-        if features is not None:
-            grouped_features = features.unsqueeze(2)
-            if self.use_xyz:
-                new_features = torch.cat(
-                    [grouped_xyz, grouped_features], dim=1
-                )  # (B, 3 + C, 1, N)
-            else:
-                new_features = grouped_features
-        else:
-            new_features = grouped_xyz
-
-        if self.ret_grouped_xyz:
-            return new_features, grouped_xyz
-        else:
-            return new_features
-
-
-class CylinderQuery(Function):
-    @staticmethod
-    def forward(ctx, radius, hmin, hmax, nsample, xyz, new_xyz, rot):
-        # type: (Any, float, float, float, int, torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor
-        r"""
-
-        Parameters
-        ----------
-        radius : float
-            radius of the cylinders
-        hmin, hmax : float
-            endpoints of cylinder height in x-rotation axis
-        nsample : int
-            maximum number of features in the cylinders
-        xyz : torch.Tensor
-            (B, N, 3) xyz coordinates of the features
-        new_xyz : torch.Tensor
-            (B, npoint, 3) centers of the cylinder query
-        rot: torch.Tensor
-            (B, npoint, 9) flatten rotation matrices from
-                           cylinder frame to world frame
-
-        Returns
-        -------
-        torch.Tensor
-            (B, npoint, nsample) tensor with the indicies of the features that form the query balls
-        """
-        return pointnet2.cylinder_query(new_xyz, xyz, rot, radius, hmin, hmax, nsample)
-
-    @staticmethod
-    def backward(ctx, a=None):
-        return None, None, None, None, None, None, None
-
-
-cylinder_query = CylinderQuery.apply
-
-
-class CylinderQueryAndGroup(nn.Module):
-    r"""
-    Groups with a cylinder query of radius and height
-
-    Parameters
-    ---------
-    radius : float32
-        Radius of cylinder
-    hmin, hmax: float32
-        endpoints of cylinder height in x-rotation axis
-    nsample : int32
-        Maximum number of features to gather in the ball
-    """
-
-    def __init__(self, radius, hmin, hmax, nsample, use_xyz=True, ret_grouped_xyz=False, normalize_xyz=False, rotate_xyz=True, sample_uniformly=False, ret_unique_cnt=False):
-        super(CylinderQueryAndGroup, self).__init__()
-        self.radius, self.nsample, self.hmin, self.hmax, = radius, nsample, hmin, hmax
-        self.use_xyz = use_xyz
-        self.ret_grouped_xyz = ret_grouped_xyz
-        self.normalize_xyz = normalize_xyz
-        self.rotate_xyz = rotate_xyz
-        self.sample_uniformly = sample_uniformly
-        self.ret_unique_cnt = ret_unique_cnt
-        if self.ret_unique_cnt:
-            assert(self.sample_uniformly)
-
-    def forward(self, xyz, new_xyz, rot, features=None):
-        r"""
-        Parameters
-        ----------
-        xyz : torch.Tensor
-            xyz coordinates of the features (B, N, 3)
-        new_xyz : torch.Tensor
-            centriods (B, npoint, 3)
-        rot : torch.Tensor
-            rotation matrices (B, npoint, 3, 3)
-        features : torch.Tensor
-            Descriptors of the features (B, C, N)
-
-        Returns
-        -------
-        new_features : torch.Tensor
-            (B, 3 + C, npoint, nsample) tensor
-        """
-        B, npoint, _ = new_xyz.size()
-        idx = cylinder_query(self.radius, self.hmin, self.hmax, self.nsample, xyz, new_xyz, rot.view(B, npoint, 9))
-
-        if self.sample_uniformly:
-            unique_cnt = torch.zeros((idx.shape[0], idx.shape[1]))
-            for i_batch in range(idx.shape[0]):
-                for i_region in range(idx.shape[1]):
-                    unique_ind = torch.unique(idx[i_batch, i_region, :])
-                    num_unique = unique_ind.shape[0]
-                    unique_cnt[i_batch, i_region] = num_unique
-                    sample_ind = torch.randint(0, num_unique, (self.nsample - num_unique,), dtype=torch.long)
-                    all_ind = torch.cat((unique_ind, unique_ind[sample_ind]))
-                    idx[i_batch, i_region, :] = all_ind
-
-
-        xyz_trans = xyz.transpose(1, 2).contiguous()
-        grouped_xyz = grouping_operation(xyz_trans, idx)  # (B, 3, npoint, nsample)
-        grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1)
-        if self.normalize_xyz:
-            grouped_xyz /= self.radius
-        if self.rotate_xyz:
-            grouped_xyz_ = grouped_xyz.permute(0, 2, 3, 1).contiguous() # (B, npoint, nsample, 3)
-            grouped_xyz_ = torch.matmul(grouped_xyz_, rot)
-            grouped_xyz = grouped_xyz_.permute(0, 3, 1, 2).contiguous()
-
-
-        if features is not None:
-            grouped_features = grouping_operation(features, idx)
-            if self.use_xyz:
-                new_features = torch.cat(
-                    [grouped_xyz, grouped_features], dim=1
-                )  # (B, C + 3, npoint, nsample)
-            else:
-                new_features = grouped_features
-        else:
-            assert (
-                self.use_xyz
-            ), "Cannot have not features and not use xyz as a feature!"
-            new_features = grouped_xyz
-
-        ret = [new_features]
-        if self.ret_grouped_xyz:
-            ret.append(grouped_xyz)
-        if self.ret_unique_cnt:
-            ret.append(unique_cnt)
-        if len(ret) == 1:
-            return ret[0]
-        else:
-            return tuple(ret)

vision_detect/vision_detect/VisionDetect/net/pointnet2/pytorch_utils.py

@@ -1,298 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-# 
-# This source code is licensed under the MIT license found in the
-# LICENSE file in the root directory of this source tree.
-
-''' Modified based on Ref: https://github.com/erikwijmans/Pointnet2_PyTorch '''
-# import torch
-import torch.nn as nn
-from typing import List, Tuple
-
-class SharedMLP(nn.Sequential):
-
-    def __init__(
-            self,
-            args: List[int],
-            *,
-            bn: bool = False,
-            activation=nn.ReLU(inplace=True),
-            preact: bool = False,
-            first: bool = False,
-            name: str = ""
-    ):
-        super().__init__()
-
-        for i in range(len(args) - 1):
-            self.add_module(
-                name + 'layer{}'.format(i),
-                Conv2d(
-                    args[i],
-                    args[i + 1],
-                    bn=(not first or not preact or (i != 0)) and bn,
-                    activation=activation
-                    if (not first or not preact or (i != 0)) else None,
-                    preact=preact
-                )
-            )
-
-
-class _BNBase(nn.Sequential):
-
-    def __init__(self, in_size, batch_norm=None, name=""):
-        super().__init__()
-        self.add_module(name + "bn", batch_norm(in_size))
-
-        nn.init.constant_(self[0].weight, 1.0)
-        nn.init.constant_(self[0].bias, 0)
-
-
-class BatchNorm1d(_BNBase):
-
-    def __init__(self, in_size: int, *, name: str = ""):
-        super().__init__(in_size, batch_norm=nn.BatchNorm1d, name=name)
-
-
-class BatchNorm2d(_BNBase):
-
-    def __init__(self, in_size: int, name: str = ""):
-        super().__init__(in_size, batch_norm=nn.BatchNorm2d, name=name)
-
-
-class BatchNorm3d(_BNBase):
-
-    def __init__(self, in_size: int, name: str = ""):
-        super().__init__(in_size, batch_norm=nn.BatchNorm3d, name=name)
-
-
-class _ConvBase(nn.Sequential):
-
-    def __init__(
-            self,
-            in_size,
-            out_size,
-            kernel_size,
-            stride,
-            padding,
-            activation,
-            bn,
-            init,
-            conv=None,
-            batch_norm=None,
-            bias=True,
-            preact=False,
-            name=""
-    ):
-        super().__init__()
-
-        bias = bias and (not bn)
-        conv_unit = conv(
-            in_size,
-            out_size,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding=padding,
-            bias=bias
-        )
-        init(conv_unit.weight)
-        if bias:
-            nn.init.constant_(conv_unit.bias, 0)
-
-        if bn:
-            if not preact:
-                bn_unit = batch_norm(out_size)
-            else:
-                bn_unit = batch_norm(in_size)
-
-        if preact:
-            if bn:
-                self.add_module(name + 'bn', bn_unit)
-
-            if activation is not None:
-                self.add_module(name + 'activation', activation)
-
-        self.add_module(name + 'conv', conv_unit)
-
-        if not preact:
-            if bn:
-                self.add_module(name + 'bn', bn_unit)
-
-            if activation is not None:
-                self.add_module(name + 'activation', activation)
-
-
-class Conv1d(_ConvBase):
-
-    def __init__(
-            self,
-            in_size: int,
-            out_size: int,
-            *,
-            kernel_size: int = 1,
-            stride: int = 1,
-            padding: int = 0,
-            activation=nn.ReLU(inplace=True),
-            bn: bool = False,
-            init=nn.init.kaiming_normal_,
-            bias: bool = True,
-            preact: bool = False,
-            name: str = ""
-    ):
-        super().__init__(
-            in_size,
-            out_size,
-            kernel_size,
-            stride,
-            padding,
-            activation,
-            bn,
-            init,
-            conv=nn.Conv1d,
-            batch_norm=BatchNorm1d,
-            bias=bias,
-            preact=preact,
-            name=name
-        )
-
-
-class Conv2d(_ConvBase):
-
-    def __init__(
-            self,
-            in_size: int,
-            out_size: int,
-            *,
-            kernel_size: Tuple[int, int] = (1, 1),
-            stride: Tuple[int, int] = (1, 1),
-            padding: Tuple[int, int] = (0, 0),
-            activation=nn.ReLU(inplace=True),
-            bn: bool = False,
-            init=nn.init.kaiming_normal_,
-            bias: bool = True,
-            preact: bool = False,
-            name: str = ""
-    ):
-        super().__init__(
-            in_size,
-            out_size,
-            kernel_size,
-            stride,
-            padding,
-            activation,
-            bn,
-            init,
-            conv=nn.Conv2d,
-            batch_norm=BatchNorm2d,
-            bias=bias,
-            preact=preact,
-            name=name
-        )
-
-
-class Conv3d(_ConvBase):
-
-    def __init__(
-            self,
-            in_size: int,
-            out_size: int,
-            *,
-            kernel_size: Tuple[int, int, int] = (1, 1, 1),
-            stride: Tuple[int, int, int] = (1, 1, 1),
-            padding: Tuple[int, int, int] = (0, 0, 0),
-            activation=nn.ReLU(inplace=True),
-            bn: bool = False,
-            init=nn.init.kaiming_normal_,
-            bias: bool = True,
-            preact: bool = False,
-            name: str = ""
-    ):
-        super().__init__(
-            in_size,
-            out_size,
-            kernel_size,
-            stride,
-            padding,
-            activation,
-            bn,
-            init,
-            conv=nn.Conv3d,
-            batch_norm=BatchNorm3d,
-            bias=bias,
-            preact=preact,
-            name=name
-        )
-
-
-class FC(nn.Sequential):
-
-    def __init__(
-            self,
-            in_size: int,
-            out_size: int,
-            *,
-            activation=nn.ReLU(inplace=True),
-            bn: bool = False,
-            init=None,
-            preact: bool = False,
-            name: str = ""
-    ):
-        super().__init__()
-
-        fc = nn.Linear(in_size, out_size, bias=not bn)
-        if init is not None:
-            init(fc.weight)
-        if not bn:
-            nn.init.constant_(fc.bias, 0)
-
-        if preact:
-            if bn:
-                self.add_module(name + 'bn', BatchNorm1d(in_size))
-
-            if activation is not None:
-                self.add_module(name + 'activation', activation)
-
-        self.add_module(name + 'fc', fc)
-
-        if not preact:
-            if bn:
-                self.add_module(name + 'bn', BatchNorm1d(out_size))
-
-            if activation is not None:
-                self.add_module(name + 'activation', activation)
-
-def set_bn_momentum_default(bn_momentum):
-
-    def fn(m):
-        if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
-            m.momentum = bn_momentum
-
-    return fn
-
-
-class BNMomentumScheduler(object):
-
-    def __init__(
-            self, model, bn_lambda, last_epoch=-1,
-            setter=set_bn_momentum_default
-    ):
-        if not isinstance(model, nn.Module):
-            raise RuntimeError(
-                "Class '{}' is not a PyTorch nn Module".format(
-                    type(model).__name__
-                )
-            )
-
-        self.model = model
-        self.setter = setter
-        self.lmbd = bn_lambda
-
-        self.step(last_epoch + 1)
-        self.last_epoch = last_epoch
-
-    def step(self, epoch=None):
-        if epoch is None:
-            epoch = self.last_epoch + 1
-
-        self.last_epoch = epoch
-        self.model.apply(self.setter(self.lmbd(epoch)))
-
-

vision_detect/vision_detect/VisionDetect/net/pointnet2/setup.py

@@ -1,61 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-# 
-# This source code is licensed under the MIT license found in the
-# LICENSE file in the root directory of this source tree.
-
-from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
-import glob
-import os
-
-
-ROOT = os.path.dirname(os.path.abspath(__file__))
-
-_ext_src_root = "_ext_src"
-
-_ext_cpp_sources = glob.glob("{}/src/*.cpp".format(_ext_src_root))
-_ext_cu_sources = glob.glob("{}/src/*.cu".format(_ext_src_root))
-
-_ext_cuda_sources = _ext_cpp_sources + _ext_cu_sources
-
-
-def is_cuda_available():
-    # return False
-    try:
-        import torch
-        return torch.cuda.is_available() and hasattr(torch.utils.cpp_extension, "CUDAExtension")
-    except (ImportError, AttributeError, RuntimeError):
-        return False
-
-_ext_modules = []
-if is_cuda_available():
-    _ext_modules.append(
-        CUDAExtension(
-            name='pointnet2.pointnet2',
-            sources=_ext_cuda_sources,
-            extra_compile_args={
-                "cxx": ["-O2", "-DWITH_CUDA", "-I{}".format("{}/{}/include".format(ROOT, _ext_src_root))],
-                "nvcc": ["-O2", "-I{}".format("{}/{}/include".format(ROOT, _ext_src_root))],
-            },
-        )
-    )
-else:
-    _ext_modules.append(
-        CppExtension(
-            name='pointnet2.pointnet2',
-            sources=_ext_cpp_sources,
-            extra_compile_args={
-                "cxx": ["-O2", "-I{}".format("{}/{}/include".format(ROOT, _ext_src_root))],
-            },
-        )
-    )
-
-
-setup(
-    name='pointnet2',
-    ext_modules=_ext_modules,
-    cmdclass={
-        'build_ext': BuildExtension
-    }
-)
-