libopenshot/src/CVTracker.cpp

/**
 * @file
 * @brief Track an object selected by the user
 * @author Jonathan Thomas <jonathan@openshot.org>
 *
 * @ref License
 */

/* LICENSE
 *
 * Copyright (c) 2008-2019 OpenShot Studios, LLC
 * <http://www.openshotstudios.com/>. This file is part of
 * OpenShot Library (libopenshot), an open-source project dedicated to
 * delivering high quality video editing and animation solutions to the
 * world. For more information visit <http://www.openshot.org/>.
 *
 * OpenShot Library (libopenshot) is free software: you can redistribute it
 * and/or modify it under the terms of the GNU Lesser General Public License
 * as published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * OpenShot Library (libopenshot) is distributed in the hope that it will be
 * useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with OpenShot Library. If not, see <http://www.gnu.org/licenses/>.
 */

#include "../include/CVTracker.h"

/// Class constructors
// Expects a tracker type, if none is passed, set default as KCF
CVTracker::CVTracker(std::string trackerType) : trackerType(trackerType){   
}
CVTracker::CVTracker() : trackerType("KCF"){
}

// Set desirable tracker method
cv::Ptr<cv::Tracker> CVTracker::select_tracker(std::string trackerType){
    cv::Ptr<cv::Tracker> t;

    if (trackerType == "BOOSTING")
        t = cv::TrackerBoosting::create();
    if (trackerType == "MIL")
        t = cv::TrackerMIL::create();
    if (trackerType == "KCF")
        t = cv::TrackerKCF::create();
    if (trackerType == "TLD")
        t = cv::TrackerTLD::create();
    if (trackerType == "MEDIANFLOW")
        t = cv::TrackerMedianFlow::create();
    if (trackerType == "GOTURN")
        t = cv::TrackerGOTURN::create();
    if (trackerType == "MOSSE")
        t = cv::TrackerMOSSE::create();
    if (trackerType == "CSRT")
        t = cv::TrackerCSRT::create();

    return t;
}

// Track object in the hole clip
void CVTracker::trackClip(openshot::Clip& video){
    // Opencv display window
    cv::namedWindow("Display Image", cv::WINDOW_NORMAL );
    // Create Tracker

    bool trackerInit = false;
    int videoLenght = video.Reader()->info.video_length;

    // Loop through video
    for (long int frame = 0; frame < videoLenght; frame++)
    {
        std::cout<<"frame: "<<frame<<"\n";
        int frame_number = frame;
        // Get current frame
        std::shared_ptr<openshot::Frame> f = video.GetFrame(frame_number);
        
        // Grab OpenCV Mat image
        cv::Mat cvimage = f->GetImageCV();

        // Pass the first frame to initialize the tracker
        if(!trackerInit){
            // Select the object to be tracked
            cv::Rect2d bbox = cv::selectROI("Display Image", cvimage);
            
            // Initialize the tracker
            initTracker(bbox, cvimage, frame_number);
            // Draw in the frame the box containing the object
            cv::rectangle(cvimage, bbox, cv::Scalar( 255, 0, 0 ), 2, 1 );

            trackerInit = true;
        }
        else{
            // Update the object tracker according to frame 
            trackerInit = trackFrame(cvimage, frame_number);
            
            // Draw box on image
            FrameData fd = GetTrackedData(frame_number);

            cv::Rect2d box(fd.x1, fd.y1, fd.x2-fd.x1, fd.y2-fd.y1);
            cv::rectangle(cvimage, box, cv::Scalar( 255, 0, 0 ), 2, 1 );
        }
        
        // Show tracking process
        cv::imshow("Display Image", cvimage);
        // Press ESC on keyboard to exit
        char c=(char)cv::waitKey(1);
        if(c==27)
            break;

    }
}

// Initialize the tracker
bool CVTracker::initTracker(cv::Rect2d initial_bbox, cv::Mat &frame, int frameId){
    // Set initial bounding box coords
    bbox = initial_bbox; 
    // Create new tracker object
    tracker = select_tracker(trackerType);

    // Initialize tracker
    tracker->init(frame, bbox);

    // Add new frame data
    trackedDataById[frameId] = FrameData(frameId, 0, bbox.x, bbox.y, bbox.x+bbox.width, bbox.y+bbox.height);

    return true;
}

// Update the object tracker according to frame 
bool CVTracker::trackFrame(cv::Mat &frame, int frameId){
    // Update the tracking result
    bool ok = tracker->update(frame, bbox);

    // Add frame number and box coords if tracker finds the object
    // Otherwise add only frame number
    if (ok)
    {
        // Add new frame data
        trackedDataById[frameId] = FrameData(frameId, 0, bbox.x, bbox.y, bbox.x+bbox.width, bbox.y+bbox.height);
    }
    else
    {
        // Add new frame data
        trackedDataById[frameId] = FrameData(frameId);
    }

    return ok;
}

bool CVTracker::SaveTrackedData(std::string outputFilePath){
    // Create tracker message
    libopenshottracker::Tracker trackerMessage;

    // Iterate over all frames data and save in protobuf message
    for(std::map<int,FrameData>::iterator it=trackedDataById.begin(); it!=trackedDataById.end(); ++it){
        FrameData fData = it->second;
        libopenshottracker::Frame* pbFrameData;
        AddFrameDataToProto(trackerMessage.add_frame(), fData);
    }

    // Add timestamp
    *trackerMessage.mutable_last_updated() = TimeUtil::SecondsToTimestamp(time(NULL));

    {
        // Write the new message to disk.
        std::fstream output(outputFilePath, ios::out | ios::trunc | ios::binary);
        if (!trackerMessage.SerializeToOstream(&output)) {
        cerr << "Failed to write protobuf message." << endl;
        return false;
        }
    }

    // Delete all global objects allocated by libprotobuf.
    google::protobuf::ShutdownProtobufLibrary();

    return true;

}

// Add frame tracked data into protobuf message.
void CVTracker::AddFrameDataToProto(libopenshottracker::Frame* pbFrameData, FrameData& fData) {

    // Save frame number and rotation
    pbFrameData->set_id(fData.frame_id);
    pbFrameData->set_rotation(0);

    libopenshottracker::Frame::Box* box = pbFrameData->mutable_bounding_box();
    // Save bounding box data
    box->set_x1(fData.x1);
    box->set_y1(fData.y1);
    box->set_x2(fData.x2);
    box->set_y2(fData.y2);
}

// Load protobuf data file
bool CVTracker::LoadTrackedData(std::string inputFilePath){
    // Create tracker message
    libopenshottracker::Tracker trackerMessage;

    {
        // Read the existing tracker message.
        fstream input(inputFilePath, ios::in | ios::binary);
        if (!trackerMessage.ParseFromIstream(&input)) {
            cerr << "Failed to parse protobuf message." << endl;
            return false;
        }
    }

    // Make sure the trackedData is empty
    trackedDataById.clear();

    // Iterate over all frames of the saved message
    for (int i = 0; i < trackerMessage.frame_size(); i++) {
        const libopenshottracker::Frame& pbFrameData = trackerMessage.frame(i);

        // Load frame and rotation data
        int id = pbFrameData.id();
        float rotation = pbFrameData.rotation();

        // Load bounding box data
        const libopenshottracker::Frame::Box& box = pbFrameData.bounding_box();
        int x1 = box.x1();
        int y1 = box.y1();
        int x2 = box.x2();
        int y2 = box.y2();

        // Assign data to tracker map
        trackedDataById[id] = FrameData(id, rotation, x1, y1, x2, y2);
    }

    // Show the time stamp from the last update in tracker data file 
    if (trackerMessage.has_last_updated()) {
        cout << "  Loaded Data. Saved Time Stamp: " << TimeUtil::ToString(trackerMessage.last_updated()) << endl;
    }

    // Delete all global objects allocated by libprotobuf.
    google::protobuf::ShutdownProtobufLibrary();

    return true;
}

// Get tracker info for the desired frame 
FrameData CVTracker::GetTrackedData(int frameId){

    // Check if the tracker info for the requested frame exists
    if ( trackedDataById.find(frameId) == trackedDataById.end() ) {
        
        return FrameData();
    } else {
        
        return trackedDataById[frameId];
    }

}