diff --git a/docker_env/gear/Dockerfile b/docker_env/gear/Dockerfile index 9186ac4..b517adf 100644 --- a/docker_env/gear/Dockerfile +++ b/docker_env/gear/Dockerfile @@ -70,4 +70,7 @@ RUN sudo apt-get update && \ ros-melodic-moveit-ros-visualization \ ros-melodic-moveit-simple-controller-manager +RUN pip3 install pyyaml +RUN pip3 install rospkg + CMD ["/bin/bash"] diff --git a/src/rrrobot_ws/src/gazebo_models/model_mappings.txt b/src/rrrobot_ws/src/gazebo_models/model_mappings.txt index 9166897..eb1d1a5 100644 --- a/src/rrrobot_ws/src/gazebo_models/model_mappings.txt +++ b/src/rrrobot_ws/src/gazebo_models/model_mappings.txt @@ -1,3 +1,3 @@ -square_trash_can,unit_box/image.jpg,plastic -unit_box,unit_box/image.jpg,plastic -plastic_bottle,plastic_bottle/image.jpg,plastic \ No newline at end of file +square_trash_can,/app/rrrobot_ws/src/gazebo_models/unit_box/image.png,plastic +unit_box,/app/rrrobot_ws/src/gazebo_models/unit_box/image.png,plastic +plastic_bottle,/app/rrrobot_ws/src/gazebo_models/plastic_bottle/images/plastic-bottle-1.jpg,plastic \ No newline at end of file diff --git a/src/rrrobot_ws/src/gazebo_models/plastic_bottle/images/plastic-bottle-1.jpg b/src/rrrobot_ws/src/gazebo_models/plastic_bottle/images/plastic-bottle-1.jpg new file mode 100644 index 0000000..7e35d58 Binary files /dev/null and b/src/rrrobot_ws/src/gazebo_models/plastic_bottle/images/plastic-bottle-1.jpg differ diff --git a/src/rrrobot_ws/src/gazebo_models/unit_box/image.png b/src/rrrobot_ws/src/gazebo_models/unit_box/image.png new file mode 100644 index 0000000..8082704 Binary files /dev/null and b/src/rrrobot_ws/src/gazebo_models/unit_box/image.png differ diff --git a/src/rrrobot_ws/src/rrrobot/src/cv_model.py b/src/rrrobot_ws/src/rrrobot/src/cv_model.py index 22b4898..b02b735 100644 --- a/src/rrrobot_ws/src/rrrobot/src/cv_model.py +++ b/src/rrrobot_ws/src/rrrobot/src/cv_model.py @@ -26,6 +26,7 @@ class Model_v1(nn.Module): def call_back(filename): # prepare input file + filename = filename.data input_image = Image.open(filename) input_image = input_image.resize((512, 384)) # print(input_image.size) @@ -49,24 +50,26 @@ def call_back(filename): print('type: ', type_dict[predicted_label]) # publish a message, name of this node is 'cv_model' - ​pub = rospy.Publisher('cv_model', String, queue_size=10) -​ ​rospy.init_node('talker', anonymous=True) -​ ​rate = rospy.Rate(10) # 10hz - - ​garbage_type = type_dict[predicted_label] - ​rospy.loginfo(filename + ':' + garbage_type) - ​pub.publish(filename + ':' + garbage_type) - ​rate.sleep() + pub = rospy.Publisher('cv_model', String, queue_size=10) + # rospy.init_node('talker', anonymous=True) + rate = rospy.Rate(10) # 10hz + + garbage_type = type_dict[predicted_label] + rospy.loginfo(filename + ':' + garbage_type) + pub.publish(filename + ':' + garbage_type) + rate.sleep() def listener(): - rospy.init_node('listener', anonymous=True) - rospy.Subscriber('/ariac/logical_camera_1', String, callback) + # rospy.init_node('listener', anonymous=True) + rospy.Subscriber('/current_image', String, call_back) # spin() simply keeps python from exiting until this node is stopped rospy.spin() if __name__ == '__main__': + rospy.init_node('cv_node') + # prepare pretrained model model = Model_v1() model.load_state_dict(torch.load('v1_0.001.pt', map_location='cpu')) diff --git a/src/rrrobot_ws/src/rrrobot/src/object_spawner_node.cpp b/src/rrrobot_ws/src/rrrobot/src/object_spawner_node.cpp index 53096bb..5eacaf2 100644 --- a/src/rrrobot_ws/src/rrrobot/src/object_spawner_node.cpp +++ b/src/rrrobot_ws/src/rrrobot/src/object_spawner_node.cpp @@ -8,6 +8,7 @@ #include #include +#include #include "rrrobot/model_insertion.h" #include @@ -25,7 +26,8 @@ public: // DEFAULT_SPAWN_POINT.x = 0.2516105; // DEFAULT_SPAWN_POINT.y = 5.474367; // DEFAULT_SPAWN_POINT.z = 0.935669; - pub = nh.advertise("/object_to_load", 1000); + model_pub = nh.advertise("/object_to_load", 1000); + image_pub = nh.advertise("/current_image", 1000); sub = nh.subscribe("/ariac/arm1/gripper/state", 1000, &ObjectSpawner::objectPlacedCallback, this); srand(time(NULL)); read_in_items(model_file); @@ -51,13 +53,21 @@ public: } private: + struct ObjectMapping + { + std::string model_name; + std::string image_file; + std::string correct_classification; + }; + const KDL::Vector DEFAULT_SPAWN_POINT; const float CONVEYOR_WIDTH; ros::NodeHandle nh; - ros::Publisher pub; + ros::Publisher model_pub; + ros::Publisher image_pub; ros::Subscriber sub; - std::vector models; // TODO: just string? or do we need - //to store the correct classification too? + std::vector models; // TODO: just string? or do we need + //to store the correct classification too? void read_in_items(const std::string &model_file) { @@ -67,10 +77,13 @@ private: { // cout << "cur: " << cur << " cur.find(,): " << cur.find(",") << endl; // line format: model_name, image_file, model_classification - std::string model_name = cur.substr(0, cur.find(",")); - // cur.erase(0, s.find(",") + 1); - // cur.erase(0, s.find(",") + 1); - models.push_back(model_name); + ObjectMapping cur_model; + cur_model.model_name = cur.substr(0, cur.find(",")); + cur.erase(0, cur.find(",") + 1); + cur_model.image_file = cur.substr(0, cur.find(",")); + cur.erase(0, cur.find(",") + 1); + cur_model.correct_classification = cur; + models.push_back(cur_model); } } @@ -79,9 +92,11 @@ private: if (ros::ok()) { rrrobot::model_insertion msg; + std_msgs::String image_file_msg; int rand_idx = rand() % models.size(); - msg.model_name = models[rand_idx]; + msg.model_name = models[rand_idx].model_name; + image_file_msg.data = models[rand_idx].image_file; cout << "Spawning " << msg.model_name << endl; // TODO: Determine position to spawn items @@ -94,7 +109,8 @@ private: msg.pose.position.y = DEFAULT_SPAWN_POINT.y(); msg.pose.position.z = DEFAULT_SPAWN_POINT.z() + 1.0; - pub.publish(msg); + model_pub.publish(msg); + image_pub.publish(image_file_msg); ros::spinOnce(); } }