dataset_creator.cpp

#include <gazebo/gazebo_client.hh>
#include <gazebo/msgs/msgs.hh>
#include <gazebo/transport/transport.hh>
#include <opencv2/opencv.hpp>
#include <opencv2/calib3d.hpp>
#include <iostream>
#include <stdlib.h>  // for strtol
#include <math.h>
#include <list>
#include <fstream>

#include "src/communication.cpp"
#include "src/vectorMod.cpp"
#include "src/geometry.cpp"
#include "src/lidarManager.cpp"

#ifndef PI
#define PI 3.14159265
#endif

static boost::mutex mutex;
const int key_left = 81;
const int key_up = 82;
const int key_down = 84;
const int key_right = 83;
const int key_esc = 27;
const int key_a = 97;
const int key_d = 100;

const float mean_vs = 0.240167305084746;
const float mean_lr = 0.731315118644068;

const int maxParam1 {300};
const int minParam1 {1};
const double incrementParam {1};
int param1 {100};
int param2 {35};
const int maxRadius {200};
const int minRadius {20};
const int maxLengthCircleVector {20};
const double f {0.01};

const float width {320.0}, height {240.0}, fov {1.047};

int main(int argc, char ** argv)
{
    // Declaration of global variables

    std::ofstream myFile("pisalau.csv");

    float speed = 0.0;
    float dir = 0.0;

    double stdDev {0.02};
    const double mean {0};
    const int kernel_size {3};
    bool show_image {true};
    float realRadius = 0.5;
    comm::point marblePoint = {5,0,0.5};
    std::vector<std::vector<cv::Vec3f>> lastCircles;
    std::vector<circleAvg> meanCenter;
    cv::Mat image, image_gray, filtered;
    comm::cameraInterface * camera (new comm::cameraInterface);
    comm::lidarInterface * lidar (new comm::lidarInterface);
    comm::poseInterface * poseRobot (new comm::poseInterface);
    comm::marbleInterface * marble (new comm::marbleInterface);

    ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    // Argc check
    if (argc > 1)
    {
        show_image = std::strtol(argv[1], nullptr,10);
    }

    ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    // Setup Gazebo and ROS Subscriptions and Publishers

    // Load Gazebo's client
    gazebo::client::setup(argc, argv);
    // Creation of node for communication
    gazebo::transport::NodePtr node (new gazebo::transport::Node());
    node->Init();

    // Create subscribers to listen to Gazebo topics
    gazebo::transport::SubscriberPtr cameraSubscriber = node->Subscribe("~/pioneer2dx/camera/link/camera/image", &comm::cameraInterface::callbackMsg, camera);
    gazebo::transport::SubscriberPtr lidarSubscriber = node->Subscribe("~/pioneer2dx/hokuyo/link/laser/scan", &comm::lidarInterface::callbackMsg, lidar);
    gazebo::transport::SubscriberPtr poseSubscriber = node->Subscribe("/gazebo/default/pose/info", &comm::poseInterface::callbackMsg, poseRobot);
    gazebo::transport::SubscriberPtr marbleSubscriber = node->Subscribe("/gazebo/default/marble/info", &comm::marbleInterface::callbackMsg, marble);

    // Publish to the robot vel_cmd topic
    gazebo::transport::PublisherPtr movementPublisher =
    node->Advertise<gazebo::msgs::Pose>("~/pioneer2dx/vel_cmd");

    // Capability to publish to reset the world
    gazebo::transport::PublisherPtr worldPublisher = node->Advertise<gazebo::msgs::WorldControl>("~/world_control");
    gazebo::msgs::WorldControl controlMessage;
    controlMessage.mutable_reset()->set_all(true);
    worldPublisher->WaitForConnection();
    worldPublisher->Publish(controlMessage);

    // TODO Create publisher to publish the location of the closest marbel
    /* std::list<std::string> listOfTopics = gazebo::transport::getAdvertisedTopics();
    std::cout << "  List of topics:" << std::endl;
    for (auto& elm: listOfTopics)
    {
        std::cout << "    -" << elm << std::endl;
    } */

    ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    // Pre-compute calibration matrixes
    float f {width/(2*tan(fov/2))};
    cv::Matx34f KA{f, 0, 0.5*320,0,
                  0, f, 0.5*240,0,
                  0, 0    , 1,0};
    cv::Matx33f K{f, 0, 0.5*320,
                  0, f, 0.5*240,
                  0, 0    , 1};
    cv::Vec<float, 5> k(-0.25, 0.12, -0.00028, 0.00005, 0); // distortion coefficients
    cv::Size frameSize(320, 240);
    cv::Mat mapX, mapY;
    cv::initUndistortRectifyMap(K, k, cv::Matx33f::eye(), K, frameSize, CV_32FC1,
                              mapX, mapY);

    ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    // Control loop
    int cntr = 0;
    while(true)
    {
        std::string savingName = "img" + std::to_string(++cntr) + ".jpg";
        //std::cout << " ---------------------------------------------------------------------------\n";

        // Declare the local variables
        cv::Mat gauss, median, blurred, sharp, laplacianGauss, laplacianBlur, laplacianMedian, laplacian, used, original, undistorted, canny, hsv;
        std::vector<cv::Vec3f> circles, circlesLap;
        comm::lidarMsg lidarFrame;

        // Make it capable to read the image 
        mutex.lock();
        int key = cv::waitKey(1);
        mutex.unlock();

        // TODO Conditions to check which key has been pressed to update the adequate values
        if (key == key_esc)
        {
            break;
        }

        //std::cout << " [NOTE] The list of possible messages is:  " << gazebo::transport::getAdvertisedTopics() << std::endl;
        //std::cout << " [NOTE] The message that is begin published is of type:  " <<  << std::endl;

        ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
        // Initial filtering of the image

        // Don't do anything if you haven't received an image
        if (!(camera->receptionAccomplished()) | !(lidar->receptionAccomplished()) | !(poseRobot->receptionAccomplished())) {continue;}

        // Obtain the messages
        original = camera->checkReceived();
        lidarFrame = lidar->checkReceived();

        // Undistort the image
        //cv::undistort(original_, original,K,k);
        cv::remap(original, undistorted, mapX, mapY, cv::INTER_LANCZOS4);

        // Create copy of original image in grayscale
        cv::cvtColor(undistorted, image_gray, cv::COLOR_BGR2GRAY);

        // Basic filterings to check which one is better
         cv::GaussianBlur(image_gray, median, cv::Size(9,9), 1.7, 1.7);

        // Segmentation/recognition of marbles
        cv::HoughCircles(median , circles, cv::HOUGH_GRADIENT, 1,
                 median.rows,
                // 13, 19, minRadius, maxRadius
                 14, 19.5, 10, 200
        );
        //std::cout << "Number of circles detected is: " << circles.size() << std::endl;

        // Filtering of the circles
        circleAvg average = averageCircle(circles);
        pushBeginning<circleAvg>(meanCenter, average, maxLengthCircleVector);
        //std::cout << "[NOTE] Real distance to marble is:        " << poseRobot->checkReceived().distance(marblePoint) - realRadius << std::endl;
        
        double distance, rdistance, ldistance, angle;

        // Generate a pose
        ignition::math::Pose3d pose(double(speed), 0, 0, 0, 0, double(dir));

        // Convert to a pose message
        gazebo::msgs::Pose msg;
        gazebo::msgs::Set(&msg, pose);
        movementPublisher->Publish(msg);

        // Calculate distance to marble
        if (circles.size() > 0)
        {
            ldistance = lidarInfo(lidarFrame, width, f, average) + mean_lr;
            distance = distanceToMarble(average, realRadius, f) + mean_vs;
            angle = angleToPointDeg(average.x_center, width, f);
            //ignition::math::Vector3<double> position = pose.Pos();
            //double pos[3] = {position.X(), position.Y(), position.Z()};
            //std::cout << pos[0] << "  --  "<< pos[1] << "  --  "<< pos[2] << std::endl;
            //rdistance = veryCustomDistance(pos[0], pos[1], pos[3], 5, 0, 0.5);
            rdistance = poseRobot->checkReceived().distance(marblePoint);
            //std::cout << "[NOTE] Calculated distance to marble is:  " << distance << std::endl;
            //std::cout << "[NOTE] Real distance to marble is:        " << poseRobot->checkReceived().distance(marblePoint) << std::endl;
            std::vector<double> calc = calcCoord(K, distance, circles[0][1], circles[0][0]);
            //std::cout << "Calculated points: x->" <<calc[0]<<"   y->"<<calc[1]<<"   z->"<<calc[2]<<std::endl;
            //std::cout << poseRobot.x_center << " -- " << marblePoint.x_center << std::endl;
            //std::cout << "";
        }
        else
        {
            distance = 0;
            rdistance = 0;
            ldistance = lidarInfo(lidarFrame, width, f, average);
        }

        if ((key == key_up) && (speed <= 1.2f))
        {
            speed += 0.05;
        }
        else if ((key == key_down) && (speed >= -1.2f))
        {
            speed -= 0.05;
        }
        else if ((key == key_right) && (dir <= 0.4f))
        {
            dir += 0.05;
        }
        else if ((key == key_left) && (dir >= -0.4f))
        {
            dir -= 0.05;
        }
        else 
        {
                speed *= 0.9;
                dir *= 0.9;
        }

        if (key == key_a)
        {
            myFile << std::to_string(circles.size()) << ",1," << std::to_string(distance) << "," << std::to_string(rdistance) << "," << std::to_string(ldistance) << "\n";
            std::cout << std::to_string(circles.size()) << ",1," << std::to_string(distance) << "," << std::to_string(rdistance) << "," << std::to_string(ldistance) << "\n";
        }

        if (key == key_d)
        {
            cntr++;
            myFile << std::to_string(circles.size()) << ",0," << std::to_string(distance) << "," << std::to_string(rdistance) << "," << std::to_string(ldistance) << "\n";
            std::cout << std::to_string(circles.size()) << ",0," << std::to_string(distance) << "," << std::to_string(rdistance) << "," << std::to_string(ldistance) << "\n";
        }

        // Obtain the real distance through obtaining the position of each object

        // Using the LIDAR to find the marbles in the range
        // Identify the angle at which the marble is.
        

        // Control of the robot based on the images


        // Show the pictures

        if (show_image){
            // Draw circles detected in each image
            cv::Mat avgCircleImage = median.clone();

            for (int i = 0; i < circles.size(); i++)
            {
                // circle center
                cv::circle( median, cv::Point(circles[i][0], circles[i][1]), 1, cv::Scalar(0,100,100), 3, cv::LINE_AA);
                // circle outline
                cv::circle( median, cv::Point(circles[i][0], circles[i][1]), circles[i][2], cv::Scalar(255,0,255), 1, cv::LINE_AA);
            }
            cv::circle(avgCircleImage, cv::Point(meanCenter[maxLengthCircleVector-1].x_center, meanCenter[maxLengthCircleVector-1].y_center), meanCenter[maxLengthCircleVector-1].radius, cv::Scalar(255,0,255), 1, cv::LINE_AA);

            /* mutex.lock();
            cv::imshow("Original", original);
            mutex.unlock();
            mutex.lock();
            cv::imshow("Undistorted", undistorted);
            mutex.unlock(); */
            mutex.lock();
            cv::imshow("Lidar", lidarFrame.im);
            mutex.unlock();

            mutex.lock();
            cv::imshow("Median", median);
            mutex.unlock();
            mutex.lock();
            cv::imshow("Average Circle", avgCircleImage);
            mutex.unlock();
        }

        //std::cout << " [NOTE] Type of the relevant message: " << gazebo::transport::getTopicMsgType("/gazebo/default/pose/info") << std::endl;
    }

    // Turn off gazebo
    gazebo::client::shutdown();

    myFile.close();

    return 0;
}