This commit is contained in:
IDONTSUDO 2024-07-03 15:25:51 +03:00
parent 559262db34
commit 50822a031d
65 changed files with 7738 additions and 1412 deletions

BIN
web_p/blender/.DS_Store vendored Normal file

Binary file not shown.

BIN
web_p/blender/assets/.DS_Store vendored Normal file

Binary file not shown.

BIN
web_p/blender/libs/.DS_Store vendored Normal file

Binary file not shown.

File diff suppressed because it is too large Load diff

File diff suppressed because it is too large Load diff

View file

@ -6,6 +6,7 @@ import blenderproc as bproc
Используется модуль blenderproc
02.05.2024 @shalenikol release 0.1
02.07.2024 @shalenikol release 0.2
"""
import numpy as np
import argparse
@ -13,25 +14,37 @@ import random
import os
import shutil
import json
from pathlib import Path
import bpy
VHACD_PATH = "blenderproc_resources/vhacd"
DIR_MODELS = "models"
DIR_MESH = "assets/libs/objects/" #"assets/mesh/"
FILE_LOG_SCENE = "res.txt"
FILE_RBS_INFO = "rbs_info.json"
FILE_GT_COCO = "scene_gt_coco.json"
EXT_MODELS = ".fbx"
TEXTURE_TMPL = "*.jpg"
Not_Categories_Name = True # наименование категории в COCO-аннотации отсутствует
def _get_list_texture(rel_path: str) -> list:
# local_path/texture/
loc = os.path.dirname(os.path.dirname(rnd_par.output_dir))
path = os.path.join(loc, rel_path)
return list(Path(path).absolute().rglob(TEXTURE_TMPL))
def _get_path_model(name_model: str) -> str:
# TODO on name_model find path for mesh (model.fbx)
# local_path/assets/mesh/
# local_path/assets/libs/objects # assets/mesh/
loc = os.path.dirname(os.path.dirname(rnd_par.output_dir))
return os.path.join(loc, "assets/mesh/"+name_model+".fbx")
return os.path.join(loc, DIR_MESH + name_model + EXT_MODELS)
def _get_path_object(name_obj: str) -> str:
# TODO on name_obj find path for scene object (object.fbx)
loc = os.path.dirname(os.path.dirname(rnd_par.output_dir))
return os.path.join(loc, "assets/mesh/"+name_obj+".fbx")
return os.path.join(loc, DIR_MESH + name_obj + EXT_MODELS)
def convert2relative(height, width, bbox):
"""
@ -43,11 +56,14 @@ def convert2relative(height, width, bbox):
return x/width, y/height, w/width, h/height
def render() -> int:
i = 0
for obj in all_meshs:
# Make the object actively participate in the physics simulation
obj.enable_rigidbody(active=True, collision_shape="COMPOUND")
# Also use convex decomposition as collision shapes
obj.build_convex_decomposition_collision_shape(VHACD_PATH)
i += 1
# print(f"{i} : {obj.get_name()}")
objs = all_meshs + rnd_par.scene.objs
@ -85,13 +101,26 @@ def render() -> int:
# Цикл рендеринга
# Do multiple times: Position the shapenet objects using the physics simulator and render X images with random camera poses
for r in range(rnd_par.n_series):
print(f"********** Series : {r+1}")
is_texture = True if "texture_path" in rnd_par.models_randomization else False
if is_texture:
val = rnd_par.models_randomization["texture_path"]
l_texture = _get_list_texture(val)
image = bpy.data.images.load(filepath=str(l_texture[r % len(l_texture)]))
# один случайный объект в кадре / все заданные объекты
random_obj = random.choice(range(rnd_par.scene.n_obj))
random_obj = random.choice(range(rnd_par.models.n_item))
meshs = []
for i,o in enumerate(all_meshs): #objs
for i,o in enumerate(all_meshs): # активные модели
if rnd_par.single_object and i != random_obj:
continue
meshs += [o]
if is_texture:
mats = o.get_materials()
for mat in mats:
# image = bpy.data.images.load(filepath=str(random.choice(l_texture)))
mat.set_principled_shader_value("Base Color", image)
for i,o in enumerate(rnd_par.scene.objs): # объекты сцены
rnd_mat = rnd_par.scene.obj_data[i]["material_randomization"]
mats = o.get_materials() #[0]
for mat in mats:
@ -99,10 +128,17 @@ def render() -> int:
mat.set_principled_shader_value("Specular", random.uniform(val[0], val[1]))
val = rnd_mat["roughness"]
mat.set_principled_shader_value("Roughness", random.uniform(val[0], val[1]))
val = rnd_mat["base_color"]
mat.set_principled_shader_value("Base Color", np.random.uniform(val[0], val[1]))
val = rnd_mat["metallic"]
mat.set_principled_shader_value("Metallic", random.uniform(val[0], val[1]))
if "texture_path" in rnd_mat: # путь к текстурам (*.jpg)
val = rnd_mat["texture_path"]
val = _get_list_texture(val)
image = bpy.data.images.load(filepath=str(random.choice(val)))
mat.set_principled_shader_value("Base Color", image)
else:
val = rnd_mat["base_color"]
mat.set_principled_shader_value("Base Color", np.random.uniform(val[0], val[1]))
# mat.set_principled_shader_value("Base Color", image)
# Randomly set the color and energy
for i,l in enumerate(ls):
@ -225,7 +261,7 @@ def render() -> int:
cat["name"] = rnd_par.models.names[i] #obj_names[i]
with open(coco_file, "w") as fh:
json.dump(data, fh, indent=0)
json.dump(data, fh, indent=1)
def explore(path: str):
if not os.path.isdir(path):
@ -331,7 +367,7 @@ if __name__ == "__main__":
models_randomization = ds_cfg["models_randomization"]
rnd_par = lambda: None
rnd_par.single_object = True
rnd_par.single_object = False # True
rnd_par.ds_name = cfg["name"]
rnd_par.output_dir = cfg["local_path"]
rnd_par.dataset_objs = cfg["dataSetObjects"]
@ -343,6 +379,7 @@ if __name__ == "__main__":
rnd_par.center_shell = cam_pos["center_shell"]
rnd_par.radius_range = cam_pos["radius_range"]
rnd_par.elevation_range = cam_pos["elevation_range"]
rnd_par.models_randomization = models_randomization
rnd_par.loc_range_low = models_randomization["loc_range_low"]
rnd_par.loc_range_high = models_randomization["loc_range_high"]
@ -353,8 +390,7 @@ if __name__ == "__main__":
bproc.init()
all_meshs = []
ret = _get_models(rnd_par, rnd_par.dataset_objs)
if ret <= 0:
if _get_models(rnd_par, rnd_par.dataset_objs) <= 0:
print("Error: no models in config")
exit(-4)
if _get_scene(rnd_par, ds_cfg["scene"]) == 0:

View file

@ -0,0 +1 @@
/home/idontsudo/robossembler_ws/src/rbs_gripper/meshes/visual/base.dae

View file

@ -0,0 +1 @@
/home/idontsudo/robossembler_ws/src/rbs_arm/meshes/visual/ee_link.dae

View file

@ -0,0 +1 @@
/home/idontsudo/robossembler_ws/src/rbs_gripper/meshes/visual/finger.dae

View file

@ -0,0 +1 @@
/home/idontsudo/robossembler_ws/src/rbs_arm/meshes/visual/fork_link.dae

View file

@ -0,0 +1 @@
/home/idontsudo/robossembler_ws/src/rbs_arm/meshes/visual/main_link.dae

View file

@ -0,0 +1,346 @@
<robot name="rbs_arm">n <link name="world" />n <joint name="base_link_joint" type="fixed">n <origin
rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0" />n <parent
link="world" />n <child link="base_link" />n </joint>n <link name="base_link">n <collision
name="base_link_collision">n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.00000 0.00000 0.00000" />n <geometry>n <mesh
filename="start_link.stl"
scale="1.00000 1.00000 1.00000" />n </geometry>n </collision> n <inertial>n <inertia
ixx="0.00503302470272442" ixy="0.000343817346410954" ixz="-4.74990755448368E-06"
iyy="0.00337962410057753" iyz="-2.3099255620051E-05" izz="0.00405858207282473" />n <origin
rpy="0.00000 0.00000 0.00000"
xyz="-0.000297002857922682 0.0964721185617698 -0.000361033370053684" /> n <mass
value="1.88031044620482" />n </inertial>n <visual name="base_link_visual">n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.00000 0.00000 0.00000" />n <geometry> n <mesh
filename="start_link.dae"
scale="1.00000 1.00000 1.00000" />n </geometry>n </visual>n </link>n <gazebo
reference="base_link">n <visual>n <material>n <diffuse>0.8 0.8 0.8 1</diffuse>n <specular>0.6
0.6
0.6 1</specular>n <ambient>1.0 1.0 1.0</ambient>n <lighting>true</lighting>n <emissive>0
0 0
1</emissive>n <pbr>n <metal>n <albedo_map>
/base_link_d.png</albedo_map> n <normal_map>
/base_link_n.png</normal_map>
n <ambient_occlusion_map>
/base_link_ao.png</ambient_occlusion_map>
n <roughness_map>
/base_link_r.png</roughness_map>
n </metal>n </pbr>n </material>n </visual>n </gazebo>n <joint name="fork0_link_joint"
type="revolute">n <limit
effort="78" lower="-6.14159" upper="6.14159" velocity="0.52" />n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.00000 0.00000 0.17833" />n <parent link="base_link" />
n <child link="fork0_link" />n <axis xyz="0 0 1" />n </joint>n <link name="fork0_link">n <collision
name="fork0_link_collision">n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.00000 -0.00000 0.00000" />n <geometry>n <mesh
filename="fork_link.stl"
scale="1.00000 1.00000 1.00000" />n </geometry>n </collision> n <inertial>n <inertia
ixx="0.00367804533572758" ixy="1.04277925521833E-05" ixz="-0.00149971410403071"
iyy="0.00415208849477534" iyz="-0.00122" izz="0.00329" />n <origin
rpy="0.00000 0.00000 0.00000"
xyz="0.0472051139085306 0.00208890925682996 0.0557265410642575" /> n <mass
value="1.12472202892859" />n </inertial>n <visual name="fork0_link_visual">n <origin
rpy="0.00000 -0.00000 -0.00000" xyz="0.00000 -0.00000 0.00000" />n <geometry> n <mesh
filename="fork_link.dae"
scale="1.00000 1.00000 1.00000" />n </geometry>n </visual>n </link>n <ros2_control
name="fork0_link_joint_hardware_interface" type="actuator">n <hardware>n <plugin>
ign_ros2_control/IgnitionSystem</plugin>n </hardware> n <joint name="fork0_link_joint">n <command_interface
name="position">n <!-- <param name="min">-1</param> -->n <!-- <param name="max">1</param> -->n </command_interface>n <command_interface
name="velocity">n <!-- <param name="min">-1</param> -->n <!-- <param name="max">1</param> -->n </command_interface>n <!-- WARN When this active a robot falls
down -->n <!-- <command_interface
name="effort"/> -->n <!-- <param name="p">${p}</param> -->n <!-- <param name="d">${d}</param> -->n <!-- <command_interface
name="effort"/> -->n <state_interface
name="position" />n <state_interface
name="velocity" />n <state_interface name="effort" />n </joint>n </ros2_control>n <gazebo
reference="fork0_link">n <visual>n <material>n <diffuse>0.8 0.8 0.8 1</diffuse>n <specular>0.6
0.6
0.6 1</specular>n <ambient>1.0 1.0 1.0</ambient>n <lighting>true</lighting>n <emissive>0
0 0
1</emissive>n <pbr>n <metal>n <albedo_map>
/fork_d.png</albedo_map>
n <normal_map>
/fork_n.png</normal_map>
n <ambient_occlusion_map>
/fork_ao.png</ambient_occlusion_map>
n <roughness_map>
/fork_r.png</roughness_map>
n </metal>n </pbr>n </material>n </visual>n </gazebo>n <joint name="main0_link_joint"
type="revolute">n <limit
effort="78" lower="-1.5708" upper="3.14159" velocity="0.52" />n <origin
rpy="-0.00000 0.00000 0.00000" xyz="0.10000 0.00000 0.09400" />n <parent
link="fork0_link" /> n <child link="main0_link" />n <axis xyz="0 1 0" />n </joint>n <link
name="main0_link">n <collision name="main0_link_collision">n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.00000 0.00000 0.00000" />n <geometry>n <mesh
filename="main_link.stl"
scale="1.00000 1.00000 1.00000" />n </geometry>n </collision> n <inertial>n <inertia
ixx="0.00315699373090845" ixy="2.84713820858537E-05" ixz="-7.01601261191721E-05"
iyy="0.00343729241263707" iyz="-0.000101485203138902" izz="0.00125534890134052" />n <origin
rpy="0.00000 0.00000 0.00000"
xyz="0.00186712264682627 -0.000412152188777604 0.0516389446895805" /> n <mass
value="1.58688811563124" />n </inertial>n <visual name="main0_link_visual">n <origin
rpy="0.00000 -0.00000 -0.00000" xyz="0.00000 0.00000 0.00000" />n <geometry> n <mesh
filename="main_link.dae"
scale="1.00000 1.00000 1.00000" />n </geometry>n </visual>n </link>n <ros2_control
name="main0_link_joint_hardware_interface" type="actuator">n <hardware>n <plugin>
ign_ros2_control/IgnitionSystem</plugin>n </hardware> n <joint name="main0_link_joint">n <command_interface
name="position">n <!-- <param name="min">-1</param> -->n <!-- <param name="max">1</param> -->n </command_interface>n <command_interface
name="velocity">n <!-- <param name="min">-1</param> -->n <!-- <param name="max">1</param> -->n </command_interface>n <!-- WARN When this active a robot falls
down -->n <!-- <command_interface
name="effort"/> -->n <!-- <param name="p">${p}</param> -->n <!-- <param name="d">${d}</param> -->n <!-- <command_interface
name="effort"/> -->n <state_interface
name="position" />n <state_interface
name="velocity" />n <state_interface name="effort" />n </joint>n </ros2_control>n <gazebo
reference="main0_link">n <visual>n <material>n <diffuse>0.8 0.8 0.8 1</diffuse>n <specular>0.6
0.6
0.6 1</specular>n <ambient>1.0 1.0 1.0</ambient>n <lighting>true</lighting>n <emissive>0
0 0
1</emissive>n <pbr>n <metal>n <albedo_map>
/main_link_d.png</albedo_map>
n <normal_map>
/main_link_n.png</normal_map>
n <ambient_occlusion_map>
/main_link_ao.png</ambient_occlusion_map>
n <roughness_map>
/main_link_r.png</roughness_map>
n </metal>n </pbr>n </material>n </visual>n </gazebo>n <joint name="fork1_link_joint"
type="revolute">n <limit
effort="78" lower="-6.14159" upper="6.14159" velocity="0.52" />n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.00000 0.00000 0.13300" />n <parent
link="main0_link" /> n <child link="fork1_link" />n <axis xyz="0 0 1" />n </joint>n <link
name="fork1_link">n <collision name="fork1_link_collision">n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.00000 -0.00000 0.00000" />n <geometry>n <mesh
filename="fork_link.stl"
scale="1.00000 1.00000 1.00000" />n </geometry>n </collision> n <inertial>n <inertia
ixx="0.00367804533572758" ixy="1.04277925521833E-05" ixz="-0.00149971410403071"
iyy="0.00415208849477534" iyz="-0.00122" izz="0.00329" />n <origin
rpy="0.00000 0.00000 0.00000"
xyz="0.0472051139085306 0.00208890925682996 0.0557265410642575" /> n <mass
value="1.12472202892859" />n </inertial>n <visual name="fork1_link_visual">n <origin
rpy="0.00000 -0.00000 -0.00000" xyz="0.00000 -0.00000 0.00000" />n <geometry> n <mesh
filename="fork_link.dae"
scale="1.00000 1.00000 1.00000" />n </geometry>n </visual>n </link>n <ros2_control
name="fork1_link_joint_hardware_interface" type="actuator">n <hardware>n <plugin>
ign_ros2_control/IgnitionSystem</plugin>n </hardware> n <joint name="fork1_link_joint">n <command_interface
name="position">n <!-- <param name="min">-1</param> -->n <!-- <param name="max">1</param> -->n </command_interface>n <command_interface
name="velocity">n <!-- <param name="min">-1</param> -->n <!-- <param name="max">1</param> -->n </command_interface>n <!-- WARN When this active a robot falls
down -->n <!-- <command_interface
name="effort"/> -->n <!-- <param name="p">${p}</param> -->n <!-- <param name="d">${d}</param> -->n <!-- <command_interface
name="effort"/> -->n <state_interface
name="position" />n <state_interface
name="velocity" />n <state_interface name="effort" />n </joint>n </ros2_control>n <gazebo
reference="fork1_link">n <visual>n <material>n <diffuse>0.8 0.8 0.8 1</diffuse>n <specular>0.6
0.6
0.6 1</specular>n <ambient>1.0 1.0 1.0</ambient>n <lighting>true</lighting>n <emissive>0
0 0
1</emissive>n <pbr>n <metal>n <albedo_map>
/fork_d.png</albedo_map>
n <normal_map>
/fork_n.png</normal_map>
n <ambient_occlusion_map>
/fork_ao.png</ambient_occlusion_map>
n <roughness_map>
/fork_r.png</roughness_map>
n </metal>n </pbr>n </material>n </visual>n </gazebo>n <joint name="main1_link_joint"
type="revolute">n <limit
effort="78" lower="-1.5708" upper="3.14159" velocity="0.52" />n <origin
rpy="-0.00000 0.00000 0.00000" xyz="0.10000 0.00000 0.09400" />n <parent
link="fork1_link" /> n <child link="main1_link" />n <axis xyz="0 1 0" />n </joint>n <link
name="main1_link">n <collision name="main1_link_collision">n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.00000 0.00000 0.00000" />n <geometry>n <mesh
filename="main_link.stl"
scale="1.00000 1.00000 1.00000" />n </geometry>n </collision> n <inertial>n <inertia
ixx="0.00315699373090845" ixy="2.84713820858537E-05" ixz="-7.01601261191721E-05"
iyy="0.00343729241263707" iyz="-0.000101485203138902" izz="0.00125534890134052" />n <origin
rpy="0.00000 0.00000 0.00000"
xyz="0.00186712264682627 -0.000412152188777604 0.0516389446895805" /> n <mass
value="1.58688811563124" />n </inertial>n <visual name="main1_link_visual">n <origin
rpy="0.00000 -0.00000 -0.00000" xyz="0.00000 0.00000 0.00000" />n <geometry> n <mesh
filename="main_link.dae"
scale="1.00000 1.00000 1.00000" />n </geometry>n </visual>n </link>n <ros2_control
name="main1_link_joint_hardware_interface" type="actuator">n <hardware>n <plugin>
ign_ros2_control/IgnitionSystem</plugin>n </hardware> n <joint name="main1_link_joint">n <command_interface
name="position">n <!-- <param name="min">-1</param> -->n <!-- <param name="max">1</param> -->n </command_interface>n <command_interface
name="velocity">n <!-- <param name="min">-1</param> -->n <!-- <param name="max">1</param> -->n </command_interface>n <!-- WARN When this active a robot falls
down -->n <!-- <command_interface
name="effort"/> -->n <!-- <param name="p">${p}</param> -->n <!-- <param name="d">${d}</param> -->n <!-- <command_interface
name="effort"/> -->n <state_interface
name="position" />n <state_interface
name="velocity" />n <state_interface name="effort" />n </joint>n </ros2_control>n <gazebo
reference="main1_link">n <visual>n <material>n <diffuse>0.8 0.8 0.8 1</diffuse>n <specular>0.6
0.6
0.6 1</specular>n <ambient>1.0 1.0 1.0</ambient>n <lighting>true</lighting>n <emissive>0
0 0
1</emissive>n <pbr>n <metal>n <albedo_map>
/main_link_d.png</albedo_map>
n <normal_map>
/main_link_n.png</normal_map>
n <ambient_occlusion_map>
/main_link_ao.png</ambient_occlusion_map>
n <roughness_map>
/main_link_r.png</roughness_map>
n </metal>n </pbr>n </material>n </visual>n </gazebo>n <joint name="fork2_link_joint"
type="revolute">n <limit
effort="78" lower="-6.14159" upper="6.14159" velocity="0.52" />n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.00000 0.00000 0.13300" />n <parent
link="main1_link" /> n <child link="fork2_link" />n <axis xyz="0 0 1" />n </joint>n <link
name="fork2_link">n <collision name="fork2_link_collision">n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.00000 -0.00000 0.00000" />n <geometry>n <mesh
filename="fork_link.stl"
scale="1.00000 1.00000 1.00000" />n </geometry>n </collision> n <inertial>n <inertia
ixx="0.00367804533572758" ixy="1.04277925521833E-05" ixz="-0.00149971410403071"
iyy="0.00415208849477534" iyz="-0.00122" izz="0.00329" />n <origin
rpy="0.00000 0.00000 0.00000"
xyz="0.0472051139085306 0.00208890925682996 0.0557265410642575" /> n <mass
value="1.12472202892859" />n </inertial>n <visual name="fork2_link_visual">n <origin
rpy="0.00000 -0.00000 -0.00000" xyz="0.00000 -0.00000 0.00000" />n <geometry> n <mesh
filename="fork_link.dae"
scale="1.00000 1.00000 1.00000" />n </geometry>n </visual>n </link>n <ros2_control
name="fork2_link_joint_hardware_interface" type="actuator">n <hardware>n <plugin>
ign_ros2_control/IgnitionSystem</plugin>n </hardware> n <joint name="fork2_link_joint">n <command_interface
name="position">n <!-- <param name="min">-1</param> -->n <!-- <param name="max">1</param> -->n </command_interface>n <command_interface
name="velocity">n <!-- <param name="min">-1</param> -->n <!-- <param name="max">1</param> -->n </command_interface>n <!-- WARN When this active a robot falls
down -->n <!-- <command_interface
name="effort"/> -->n <!-- <param name="p">${p}</param> -->n <!-- <param name="d">${d}</param> -->n <!-- <command_interface
name="effort"/> -->n <state_interface
name="position" />n <state_interface
name="velocity" />n <state_interface name="effort" />n </joint>n </ros2_control>n <gazebo
reference="fork2_link">n <visual>n <material>n <diffuse>0.8 0.8 0.8 1</diffuse>n <specular>0.6
0.6
0.6 1</specular>n <ambient>1.0 1.0 1.0</ambient>n <lighting>true</lighting>n <emissive>0
0 0
1</emissive>n <pbr>n <metal>n <albedo_map>
/fork_d.png</albedo_map>
n <normal_map>
/fork_n.png</normal_map>
n <ambient_occlusion_map>
/fork_ao.png</ambient_occlusion_map>
n <roughness_map>
/fork_r.png</roughness_map>
n </metal>n </pbr>n </material>n </visual>n </gazebo>n <joint name="ee_link_joint"
type="revolute">n <limit
effort="78" lower="-1.5708" upper="3.14159" velocity="0.52" />n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.10000 0.00000 0.09473" />n <parent
link="fork2_link" /> n <child link="ee_link" />n <axis xyz="0 1 0" />n </joint>n <joint
name="tool0_joint" type="fixed">n <origin rpy="0.00000 0.00000 0.00000"
xyz="0.00000 0.00000 0.11000" />n <parent link="ee_link" />n <child link="tool0" />n </joint>
n <link name="tool0" />n <link
name="ee_link">n <collision name="ee_link_collision">n <origin
rpy="-0.00000 0.00000 -0.00000" xyz="0.00000 -0.00000 0.00000" />n <geometry>n <mesh
filename="ee_link.stl"
scale="1.00000 1.00000 1.00000" />n </geometry>n </collision>n <inertial>n <inertia
ixx="0.00147695259043549" ixy="-2.66894744420299E-05"
ixz="-4.40871314563273E-05" iyy="0.00135500487881796" iyz="-3.19001462979333E-05"
izz="0.00087582892706912" />n <origin rpy="0.00000 0.00000 0.00000"
xyz="-9.7531539777207E-06 -0.000888494418875867 0.0342332199538358" />n <mass
value="1.88031044620482" />n </inertial> n <visual name="ee_link_visual">n <origin
rpy="0.00000 0.00000 0.00000" xyz="0.00000 -0.00000 0.00000" />n <geometry>n <mesh
filename="ee_link.dae"
scale="1.00000 1.00000 1.00000" />n </geometry> n </visual>n </link>n <ros2_control
name="ee_link_joint_hardware_interface" type="actuator">n <hardware>n <plugin>
ign_ros2_control/IgnitionSystem</plugin>n </hardware>n <joint name="ee_link_joint">n <command_interface
name="position">n <!-- <param name="min">-1</param> -->n <!-- <param name="max">1</param> -->n </command_interface>n <command_interface name="velocity">n <!--
<param
name="min">-1</param> -->n <!--
<param
name="max">1</param> -->n </command_interface>
n <!-- WARN When this active a robot falls down -->n <!-- <command_interface name="effort"/> -->n <!-- <param name="p">${p}</param> -->n <!-- <param name="d">${d}</param> -->n <!-- <command_interface name="effort"/> -->n <state_interface name="position" />n <state_interface name="velocity" />n <state_interface
name="effort" />n </joint>n </ros2_control>n <gazebo reference="ee_link">n <visual>n <material>
n <diffuse>0.8 0.8 0.8 1</diffuse>n <specular>0.6 0.6 0.6 1</specular>n <ambient>1.0
1.0 1.0</ambient>n <lighting>true</lighting>n <emissive>0 0 0 1</emissive>n <pbr>
n <metal>n <albedo_map>
/ee_link_d.png</albedo_map>
n <normal_map>
/ee_link_n.png</normal_map>
n <ambient_occlusion_map>
/ee_link_ao.png</ambient_occlusion_map>
n <roughness_map>
/ee_link_r.png</roughness_map>
n </metal>n </pbr>n </material>n </visual>n </gazebo>n <!-- END robot description -->n <ros2_control name="rbs_gripper_gazebo"
type="system">n <!-- Plugins -->n <hardware>n <plugin>ign_ros2_control/IgnitionSystem</plugin>n </hardware>n <!--
Joint interfaces -->
n <joint name="rbs_gripper_rot_base_joint">n <state_interface name="position">n <param
name="initial_value">0.000</param>n </state_interface> n <command_interface
name="position" />n <command_interface name="velocity" />n <command_interface
name="effort" />n <state_interface name="velocity" />n <state_interface
name="effort" />n </joint>n <joint
name="rbs_gripper_r_finger_joint">n <state_interface name="position">n <param
name="initial_value">0.000</param>n </state_interface> n <command_interface
name="position" />n <command_interface name="velocity" />n <command_interface
name="effort" />n <state_interface name="velocity" />n <state_interface
name="effort" />n </joint>n <joint
name="rbs_gripper_l_finger_joint">n <param name="mimic">rbs_gripper_r_finger_joint</param>
n <param name="multiplier">1</param> n </joint>n </ros2_control>n <joint
name="rbs_gripper_gripper_base_joint" type="fixed">n <origin rpy="0 0 0" xyz="0 0 0" />n <parent
link="tool0" />n <child link="rbs_gripper_gripper_base_link" />n </joint>n <link
name="rbs_gripper_gripper_base_link">n <inertial>n <origin rpy="0 0 0"
xyz="0.000364704367134063 0.0336387482840125 0.0593891203954369" />n <mass
value="1.13983632906086" />n <inertia ixx="0.00107738806534129"
ixy="-1.09841172461737E-05" ixz="2.62750043451545E-06" iyy="0.000717388573992299"
iyz="-2.95426438182787E-05" izz="0.00115777179755934" /> n </inertial>n <visual>n <origin
rpy="0 0 0" xyz="0 0 0" />n <geometry>n <mesh
filename="base.dae" /> n </geometry> n <material name="rbs_gripper_material">n <color
rgba="0.752941176470588 0.752941176470588 0.752941176470588 1" />n </material>n </visual>
n <collision>n <origin
rpy="0 0 0" xyz="0 0 0" />n <geometry>n <mesh
filename="base.stl" /> n </geometry>n </collision>n </link> n <link
name="rbs_gripper_rot_base_link">n <inertial>n <origin
rpy="0 0 0" xyz="6.79110135283868E-11 -3.80956832067611E-10 0.00775394473595793" />n <mass
value="0.161003401535982" />n <inertia ixx="0.00011089089949771"
ixy="5.01335040610636E-06" ixz="1.74608448389267E-14" iyy="0.000105515893695012"
iyz="-2.03282362854432E-14" izz="0.000206912001661452" />n </inertial>n <visual>n <origin
rpy="0 0 0"
xyz="0 0 0" />n <geometry>n <mesh
filename="rotor.dae" /> n </geometry>n <material name="rbs_gripper_material">n <color
rgba="0.250980392156863 0.250980392156863 0.250980392156863 1" />n </material>n </visual>
n <collision>n <origin rpy="0 0 0" xyz="0 0 0" />n <geometry> n <mesh
filename="rotor.stl" /> n </geometry>n </collision>n </link>n <joint
name="rbs_gripper_rot_base_joint" type="revolute"> n <origin
rpy="0 0 0" xyz="0 0 0.10861" />n <parent link="rbs_gripper_gripper_base_link" />n <child
link="rbs_gripper_rot_base_link" />n <axis xyz="0 0 1" />n <limit
effort="78" lower="-3.14159" upper="3.14159" velocity="0.52" />n </joint>n <link
name="rbs_gripper_l_finger_link">n <inertial>n <origin
rpy="0 0 0" xyz="0.00399878118534129 0.0187296413885176 -0.0777776233934166" />n <mass
value="0.0601996441483964" />n <inertia ixx="4.18533281165612E-05"
ixy="-7.11657995951147E-06" ixz="1.81029223490065E-05" iyy="7.89886367868258E-05"
iyz="1.20542845942065E-05"
izz="5.16740841307935E-05" />n </inertial>n <visual>n <origin rpy="0 0 0"
xyz="0 0 0" />n <geometry>n <mesh
filename="finger.dae" /> n </geometry>n <material name="rbs_gripper_material">n <color
rgba="0.752941176470588 0 0 1" />n </material> n </visual>n <collision>n <origin
rpy="0 0 0" xyz="0 0 0" />n <geometry>n <mesh
filename="finger.stl" /> n </geometry> n </collision>n </link>n <joint
name="rbs_gripper_l_finger_joint" type="prismatic"> n <origin rpy="0 0 1.5708"
xyz="0 0 0.1071" />n <parent
link="rbs_gripper_rot_base_link" />n <child link="rbs_gripper_l_finger_link" />n <axis
xyz="-1 0 0" />n <limit effort="10" lower="0" upper="0.064"
velocity="0.53" />n <mimic joint="rbs_gripper_r_finger_joint" multiplier="1" />n </joint>
n <link name="rbs_gripper_r_finger_link">n <inertial> n <origin rpy="0 0 0"
xyz="0.0039988 -0.077778 -0.01873" />n <mass value="0.0602" />n <inertia
ixx="4.1853E-05" ixy="1.8103E-05" ixz="7.1166E-06" iyy="5.1674E-05"
iyz="-1.2054E-05" izz="7.8989E-05" />n </inertial>n <visual>n <origin rpy="0 0 0"
xyz="0 0 0" />n <geometry>n <mesh
filename="finger.dae" /> n </geometry>n <material name="rbs_gripper_material">n <color
rgba="0.75294 0 0 1" />n </material> n </visual>n <collision>n <origin
rpy="0 0 0" xyz="0 0 0" />n <geometry>n <mesh
filename="finger.stl" /> n </geometry> n </collision>n </link>n <joint
name="rbs_gripper_r_finger_joint" type="prismatic"> n <origin rpy="0 0 -1.5708"
xyz="0 0 0.1071" />n <parent
link="rbs_gripper_rot_base_link" />n <child link="rbs_gripper_r_finger_link" />n <axis
xyz="-1 0 0" />n <limit effort="10" lower="0" upper="0.064"
velocity="0.53" />n </joint>n <link name="gripper_grasp_point" />n <joint
name="rbs_gripper_gripper_tool0_joint" type="fixed">n <origin rpy="0 0 0"
xyz="0 0 0.09139" />n <parent link="rbs_gripper_rot_base_link" />n <child
link="gripper_grasp_point" />n </joint>n <ros2_control name="fts_sensor"
type="sensor">n <hardware>n <plugin>ign_ros2_control/IgnitionFts</plugin>n </hardware>n <sensor
name="fts_sensor">n <state_interface name="force.x" />n <state_interface name="force.y" />
n <state_interface
name="force.z" />n <state_interface name="torque.x" />n <state_interface
name="torque.y" />n <state_interface
name="torque.z" />n </sensor>n </ros2_control>n <gazebo reference="tool0_joint">n <preserveFixedJoint>
true</preserveFixedJoint>n <sensor name="fts_sensor" type="force_torque">n <always_on>true</always_on>
n <update_rate>
50</update_rate>n <visualize>true</visualize>n <topic>ft_data</topic>n <force_torque>
n <frame>
sensor</frame>n <measure_direction>child_to_parent</measure_direction>n </force_torque>
n </sensor> n </gazebo>n <gazebo>n <plugin filename="libign_ros2_control-system.so"
name="ign_ros2_control::IgnitionROS2ControlPlugin">n <parameters>
/home/idontsudo/robossembler_ws/install/rbs_arm/share/rbs_arm/config/rbs_arm0_controllers.yaml</parameters>
n <ros>n <namespace></namespace>n </ros>n </plugin>n </gazebo>n</robot>

View file

@ -0,0 +1 @@
/home/idontsudo/robossembler_ws/src/rbs_gripper/meshes/visual/rotor.dae

View file

@ -0,0 +1 @@
/home/idontsudo/robossembler_ws/src/rbs_arm/meshes/visual/start_link.dae

View file

@ -1,20 +1,531 @@
"""
train_Dope
train_Dope
Общая задача: оценка позиции объекта (Pose estimation)
Реализуемая функция: обучение нейросетевой модели DOPE по заданному BOP-датасету
python3 $PYTHON_EDUCATION --path /Users/idontsudo/webservice/server/build/public/7065d6b6-c8a3-48c5-9679-bb8f3a690296 \
python3 $PYTHON_EDUCATION --path /Users/user/webservice/server/build/public/7065d6b6-c8a3-48c5-9679-bb8f3a690296 \
--name test1234 --datasetName 32123213
25.04.2024 @shalenikol release 0.1
08.05.2024 @shalenikol release 0.1
"""
import os
import json
import shutil
import numpy as np
import transforms3d as t3d
def train_Dope_i(path:str, wname:str, dname:str, outpath:str, epochs:int):
results = f"torchrun --nproc_per_node=1 train.py --local_rank 0 --data {os.path.join(path,dname)} --object fork" \
+ f" -e {epochs} --batchsize 16 --exts jpg --imagesize 640 --pretrained" \
+ " --net_path /home/shalenikol/fork_work/dope_training/output/weights_2996/net_epoch_47.pth"
print(results)
FILE_RBS_INFO = "rbs_info.json"
FILE_CAMERA = "camera.json"
FILE_GT = "scene_gt.json"
FILE_GT_COCO = "scene_gt_coco.json"
FILE_GT_INFO = "scene_gt_info.json"
FILE_MODEL = "epoch"
EXT_MODEL = ".pth"
EXT_RGB = "jpg"
DIR_ROOT_DS = "dataset_dope"
DIR_TRAIN_OUT = "out_weights"
MODEL_SCALE = 1000 # исходная модель в метрах, преобразуем в мм (для DOPE)
# Own_Numbering_Files = True # наименование image-файлов: собственная нумерация
nn_image = 0
K_intrinsic = []
model_info = []
camera_data = {}
im_width = 0
nb_update_network = 0
# [
# [min(x), min(y), min(z)],
# [min(x), max(y), min(z)],
# [min(x), max(y), max(z)],
# [min(x), min(y), max(z)],
# [max(x), min(y), max(z)],
# [max(x), max(y), min(z)],
# [max(x), max(y), max(z)],
# [max(x), min(y), max(z)],
# [xc, yc, zc] # min + (max - min) / 2
# ]
def trans_3Dto2D_point_in_camera(xyz, K_m, R_m2c, t_m2c):
"""
xyz : 3D-координаты точки
K_m : внутренняя матрица камеры 3х3
R_m2c : матрица поворота 3х3
t_m2c : вектор перемещения 3х1
return [u,v]
"""
K = np.array(K_m)
r = np.array(R_m2c)
r.shape = (3, 3)
t = np.array(t_m2c)
t.shape = (3, 1)
T = np.concatenate((r, t), axis=1)
P_m = np.array(xyz)
P_m.resize(4)
P_m[-1] = 1.0
P_m.shape = (4, 1)
# Project (X, Y, Z, 1) into cameras coordinate system
P_c = T @ P_m # 4x1
# Apply camera intrinsics to map (Xc, Yc, Zc) to p=(x, y, z)
p = K @ P_c
# Normalize by z to get (u,v,1)
uv = (p / p[2][0])[:-1]
return uv.flatten().tolist()
def gt_parse(path: str, out_dir: str):
global nn_image
with open(os.path.join(path, FILE_GT_COCO), "r") as fh:
coco_data = json.load(fh)
with open(os.path.join(path, FILE_GT), "r") as fh:
gt_data = json.load(fh)
with open(os.path.join(path, FILE_GT_INFO), "r") as fh:
gt_info = json.load(fh)
for img in coco_data["images"]:
rgb_file = os.path.join(path, img["file_name"])
if os.path.isfile(rgb_file):
# if Own_Numbering_Files:
ext = os.path.splitext(rgb_file)[1] # only ext
f = f"{nn_image:06}"
out_img = os.path.join(out_dir, f + ext)
# else:
# f = os.path.split(rgb_file)[1] # filename with extension
# f = os.path.splitext(f)[0] # only filename
# out_img = out_dir
shutil.copy2(rgb_file, out_img)
out_file = os.path.join(out_dir,f+".json")
nn_image += 1
# full annotation of the one image
all_data = camera_data.copy()
cat_names = {obj["id"]: obj["name"] for obj in coco_data["categories"]}
id_img = img["id"] # 0, 1, 2 ...
sid_img = str(id_img) # "0", "1", "2" ...
img_info = gt_info[sid_img]
img_gt = gt_data[sid_img]
img_idx = 0 # object index on the image
objs = []
for ann in coco_data["annotations"]:
if ann["image_id"] == id_img:
item = ann["category_id"]
obj_data = {}
obj_data["class"] = cat_names[item]
x, y, width, height = ann["bbox"]
obj_data["bounding_box"] = {"top_left":[x,y], "bottom_right":[x+width,y+height]}
# visibility from FILE_GT_INFO
item_info = img_info[img_idx]
obj_data["visibility"] = item_info["visib_fract"]
# location from FILE_GT
item_gt = img_gt[img_idx]
obj_id = item_gt["obj_id"] - 1 # index with 0
cam_R_m2c = item_gt["cam_R_m2c"]
cam_t_m2c = item_gt["cam_t_m2c"]
obj_data["location"] = cam_t_m2c
q = t3d.quaternions.mat2quat(np.array(cam_R_m2c))
obj_data["quaternion_xyzw"] = [q[1], q[2], q[3], q[0]]
cuboid_xyz = model_info[obj_id]
obj_data["projected_cuboid"] = [
trans_3Dto2D_point_in_camera(cub, K_intrinsic, cam_R_m2c, cam_t_m2c)
for cub in cuboid_xyz
]
objs.append(obj_data)
img_idx += 1
all_data["objects"] = objs
with open(out_file, "w") as fh:
json.dump(all_data, fh, indent=2)
def explore(path: str, res_dir: str):
if not os.path.isdir(path):
return
folders = [
os.path.join(path, o)
for o in os.listdir(path)
if os.path.isdir(os.path.join(path, o))
]
for path_entry in folders:
if os.path.isfile(os.path.join(path_entry,FILE_GT_COCO)) and \
os.path.isfile(os.path.join(path_entry,FILE_GT_INFO)) and \
os.path.isfile(os.path.join(path_entry,FILE_GT)):
gt_parse(path_entry, res_dir)
else:
explore(path_entry, res_dir)
def BOP2DOPE_dataset(dpath: str, out_dir: str) -> str:
""" Convert BOP-dataset to YOLO format for train """
res_dir = os.path.join(out_dir, DIR_ROOT_DS)
if os.path.isdir(res_dir):
shutil.rmtree(res_dir)
os.mkdir(res_dir)
explore(dpath, res_dir)
return out_dir
def train(dopepath:str, wname:str, epochs:int, pretrain: bool, lname: list):
import random
# try:
import configparser as configparser
# except ImportError:
# import ConfigParser as configparser
import torch
# import torch.nn.parallel
import torch.optim as optim
import torch.utils.data
import torchvision.transforms as transforms
from torch.autograd import Variable
import datetime
from tensorboardX import SummaryWriter
from models_dope import DopeNetwork
from utils_dope import CleanVisiiDopeLoader #, VisualizeBeliefMap, save_image
import warnings
warnings.filterwarnings("ignore")
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
torch.autograd.set_detect_anomaly(False)
torch.autograd.profiler.profile(False)
torch.autograd.gradcheck = False
torch.backends.cudnn.benchmark = True
start_time = datetime.datetime.now()
print("start:", start_time.strftime("%m/%d/%Y, %H:%M:%S"))
res_model = os.path.join(dopepath, wname + EXT_MODEL)
local_rank = 0
opt = lambda: None
opt.use_s3 = False
opt.train_buckets = []
opt.endpoint = None
opt.lr=0.0001
opt.loginterval=100
opt.sigma=0.5 # 4
opt.nbupdates=None
# opt.save=False
# opt.option="default"
# opt.gpuids=[0]
opt.namefile=FILE_MODEL
opt.workers=8
opt.batchsize=16
opt.data = [os.path.join(dopepath, DIR_ROOT_DS)]
opt.outf = os.path.join(dopepath, DIR_TRAIN_OUT)
opt.object = lname #["fork"]
opt.exts = [EXT_RGB]
# opt.imagesize = im_width
opt.epochs = epochs
opt.pretrained = pretrain
opt.net_path = res_model if pretrain else None
opt.manualseed = random.randint(1, 10000)
# # Validate Arguments
# if opt.use_s3 and (opt.train_buckets is None or opt.endpoint is None):
# raise ValueError(
# "--train_buckets and --endpoint must be specified if training with data from s3 bucket."
# )
# if not opt.use_s3 and opt.data is None:
# raise ValueError("--data field must be specified.")
os.makedirs(opt.outf, exist_ok=True)
# if local_rank == 0:
# writer = SummaryWriter(opt.outf + "/runs/")
random.seed(opt.manualseed)
torch.cuda.set_device(local_rank)
# torch.distributed.init_process_group(backend="nccl", init_method="env://")
torch.manual_seed(opt.manualseed)
torch.cuda.manual_seed_all(opt.manualseed)
# # Data Augmentation
# if not opt.save:
# contrast = 0.2
# brightness = 0.2
# noise = 0.1
# normal_imgs = [0.59, 0.25]
# transform = transforms.Compose(
# [
# AddRandomContrast(0.2),
# AddRandomBrightness(0.2),
# transforms.Resize(opt.imagesize),
# ]
# )
# else:
# contrast = 0.00001
# brightness = 0.00001
# noise = 0.00001
# normal_imgs = None
# transform = transforms.Compose(
# [transforms.Resize(opt.imagesize), transforms.ToTensor()]
# )
# Load Model
net = DopeNetwork()
output_size = 50
# opt.sigma = 0.5
train_dataset = CleanVisiiDopeLoader(
opt.data,
sigma=opt.sigma,
output_size=output_size,
extensions=opt.exts,
objects=opt.object,
use_s3=opt.use_s3,
buckets=opt.train_buckets,
endpoint_url=opt.endpoint,
)
trainingdata = torch.utils.data.DataLoader(
train_dataset,
batch_size=opt.batchsize,
shuffle=True,
num_workers=opt.workers,
pin_memory=True,
)
if not trainingdata is None:
print(f"training data: {len(trainingdata)} batches")
print("Loading Model...")
net = net.cuda()
# net = torch.nn.parallel.DistributedDataParallel(
# net.cuda(), device_ids=[local_rank], output_device=local_rank
# )
if opt.pretrained:
if opt.net_path is not None:
net.load_state_dict(torch.load(opt.net_path))
else:
print("Error: Did not specify path to pretrained weights.")
quit()
parameters = filter(lambda p: p.requires_grad, net.parameters())
optimizer = optim.Adam(parameters, lr=opt.lr)
print("ready to train!")
global nb_update_network
nb_update_network = 0
# best_results = {"epoch": None, "passed": None, "add_mean": None, "add_std": None}
scaler = torch.cuda.amp.GradScaler()
def _runnetwork(epoch, train_loader): #, syn=False
global nb_update_network
# net
net.train()
loss_avg_to_log = {}
loss_avg_to_log["loss"] = []
loss_avg_to_log["loss_affinities"] = []
loss_avg_to_log["loss_belief"] = []
loss_avg_to_log["loss_class"] = []
for batch_idx, targets in enumerate(train_loader):
optimizer.zero_grad()
data = Variable(targets["img"].cuda())
target_belief = Variable(targets["beliefs"].cuda())
target_affinities = Variable(targets["affinities"].cuda())
output_belief, output_aff = net(data)
loss = None
loss_belief = torch.tensor(0).float().cuda()
loss_affinities = torch.tensor(0).float().cuda()
loss_class = torch.tensor(0).float().cuda()
for stage in range(len(output_aff)): # output, each belief map layers.
loss_affinities += (
(output_aff[stage] - target_affinities)
* (output_aff[stage] - target_affinities)
).mean()
loss_belief += (
(output_belief[stage] - target_belief)
* (output_belief[stage] - target_belief)
).mean()
loss = loss_affinities + loss_belief
# if batch_idx == 0:
# post = "train"
# if local_rank == 0:
# for i_output in range(1):
# # input images
# writer.add_image(
# f"{post}_input_{i_output}",
# targets["img_original"][i_output],
# epoch,
# dataformats="CWH",
# )
# # belief maps gt
# imgs = VisualizeBeliefMap(target_belief[i_output])
# img, grid = save_image(
# imgs, "some_img.png", mean=0, std=1, nrow=3, save=False
# )
# writer.add_image(
# f"{post}_belief_ground_truth_{i_output}",
# grid,
# epoch,
# dataformats="CWH",
# )
# # belief maps guess
# imgs = VisualizeBeliefMap(output_belief[-1][i_output])
# img, grid = save_image(
# imgs, "some_img.png", mean=0, std=1, nrow=3, save=False
# )
# writer.add_image(
# f"{post}_belief_guess_{i_output}",
# grid,
# epoch,
# dataformats="CWH",
# )
loss.backward()
optimizer.step()
nb_update_network += 1
# log the loss
loss_avg_to_log["loss"].append(loss.item())
loss_avg_to_log["loss_class"].append(loss_class.item())
loss_avg_to_log["loss_affinities"].append(loss_affinities.item())
loss_avg_to_log["loss_belief"].append(loss_belief.item())
if batch_idx % opt.loginterval == 0:
print(
"Train Epoch: {} [{}/{} ({:.0f}%)] \tLoss: {:.15f} \tLocal Rank: {}".format(
epoch,
batch_idx * len(data),
len(train_loader.dataset),
100.0 * batch_idx / len(train_loader),
loss.item(),
local_rank,
)
)
# # log the loss values
# if local_rank == 0:
# writer.add_scalar("loss/train_loss", np.mean(loss_avg_to_log["loss"]), epoch)
# writer.add_scalar("loss/train_cls", np.mean(loss_avg_to_log["loss_class"]), epoch)
# writer.add_scalar("loss/train_aff", np.mean(loss_avg_to_log["loss_affinities"]), epoch)
# writer.add_scalar("loss/train_bel", np.mean(loss_avg_to_log["loss_belief"]), epoch)
for epoch in range(1, opt.epochs + 1):
_runnetwork(epoch, trainingdata)
try:
if local_rank == 0:
torch.save(
net.state_dict(),
f"{opt.outf}/{opt.namefile}_{str(epoch).zfill(3)}.pth",
)
except Exception as e:
print(f"Encountered Exception: {e}")
if not opt.nbupdates is None and nb_update_network > int(opt.nbupdates):
break
# if local_rank == 0:
# save result model
torch.save(net.state_dict(), res_model) #os.path.join(dopepath, wname + EXT_MODEL))
# else:
# torch.save(
# net.state_dict(),
# f"{opt.outf}/{opt.namefile}_{str(epoch).zfill(3)}_rank_{local_rank}.pth",
# )
print("end:", datetime.datetime.now().strftime("%m/%d/%Y, %H:%M:%S"))
print("Total time taken: ", str(datetime.datetime.now() - start_time).split(".")[0])
def train_Dope_i(path:str, wname:str, dname:str, outpath:str, epochs:int, pretrain: bool):
""" Main procedure for train DOPE model """
global K_intrinsic, model_info, camera_data, im_width
if not os.path.isdir(outpath):
print(f"Invalid output path '{outpath}'")
exit(-1)
out_dir = os.path.join(outpath, wname)
ds_path = os.path.join(path, dname)
if not os.path.isdir(ds_path):
print(f"{ds_path} : no BOP directory")
return ""
camera_json = os.path.join(ds_path, FILE_CAMERA)
if not os.path.isfile(camera_json):
print(f"{camera_json} : no intrinsic camera file")
return ""
rbs_info = os.path.join(ds_path, FILE_RBS_INFO)
if not os.path.isfile(rbs_info):
print(f"{rbs_info} : no dataset info file")
return ""
camera_data = {}
with open(camera_json, "r") as fh:
data = json.load(fh)
keys = ["cx","cy","fx","fy"]
intrinsic = {k: data[k] for k in keys}
im_height = data["height"]
im_width = data["width"]
camera_data["camera_data"] = dict(intrinsic=intrinsic, height=im_height, width=im_width)
K_intrinsic = [
[data["fx"], 0.0, data["cx"]],
[0.0, data["fy"], data["cy"]],
[0.0, 0.0, 1.0]
]
# calc cuboid + center
with open(rbs_info, "r") as fh:
info = json.load(fh)
# список имён объектов
list_name = list(map(lambda x: x["name"], info))
# in FILE_RBS_INFO model numbering from smallest to largest
model_info = []
for m_info in info:
cub = np.array(m_info["cuboid"]) * MODEL_SCALE
xyz_min = cub.min(axis=0)
xyz_max = cub.max(axis=0)
# [xc, yc, zc] # min + (max - min) / 2
center = []
for i in range(3):
center.append(xyz_min[i] + (xyz_max[i]- xyz_min[i]) / 2)
c = np.array(center, ndmin=2)
model_info.append(np.append(cub, c, axis=0))
if pretrain:
# продолжить обучение
if not os.path.isdir(out_dir):
print(f"No dir '{out_dir}'")
exit(-2)
dpath = out_dir
# model_path = os.path.join(dpath, wname + ".pt")
else:
# обучение сначала
if not os.path.isdir(out_dir):
os.mkdir(out_dir)
dpath = BOP2DOPE_dataset(ds_path, out_dir)
if len(dpath) == 0:
print(f"Error in convert dataset '{ds_path}' to '{outpath}'")
exit(-4)
# model_path = os.path.join(dpath, FILE_BASEMODEL)
# results = f"python train.py --local_rank 0 --data {dpath} --object fork" \
# + f" -e {epochs} --batchsize 16 --exts jpg --imagesize 640 --pretrained" \
# + " --net_path /home/shalenikol/fork_work/dope_training/output/weights_2996/net_epoch_47.pth"
# print(results)
train(dpath, wname, epochs, pretrain, list_name)
import argparse
@ -25,6 +536,7 @@ if __name__ == "__main__":
parser.add_argument("--datasetName", required=True, help="String with dataset name")
parser.add_argument("--outpath", default="weights", help="Output path for weights")
parser.add_argument("--epoch", default=3, help="How many training epochs")
parser.add_argument('--pretrain', action="store_true", help="Use pretraining")
args = parser.parse_args()
train_Dope_i(args.path, args.name, args.datasetName, args.outpath, args.epoch)
train_Dope_i(args.path, args.name, args.datasetName, args.outpath, args.epoch, args.pretrain)

View file

@ -28,7 +28,7 @@ DIR_COCO_DS = "rbs_coco"
DIR_RGB_DS = "images"
DIR_LABELS_DS = "labels"
SZ_SERIES = 5 # number of train images per validation images
SZ_SERIES = 15 # number of train images per validation images
nn_image = 0
f1 = f2 = None
@ -166,7 +166,6 @@ def train_YoloV8(path:str, wname:str, dname:str, outpath:str, epochs:int, pretra
shutil.copy2(wf, os.path.join(dpath, wname + ".pt"))
shutil.rmtree(results.save_dir)
# print(f"\n ********\n{wf}")
if __name__ == "__main__":
import argparse