DeskSweeper: A robot that cleans your desk automatically

/*
ArmController

- SG90
- 50Hz -> 20msec
- -90deg: 0.5msec = 6.4 / 256
- 0deg: 1.45msec = 18.56 / 256
- 90deg: 2.4msec = 30.72 / 256

*/
#include <M5Unified.h>

class ArmController {
public:
    static const uint16_t MAX_ACTIONS = 1000;
    double PWM_FREQUENCY = 50;
    double PWM_RESOLUTION = 8;
    uint8_t SHOULDER_CHANNEL = 0;
    uint8_t ELBOW_CHANNEL = 1;
    uint8_t GRIPPER_CHANNEL = 2;

    int getActionIndex() const { return _action_index; }
    int getActionCount() const { return _action_count; }
    int getTargetShoulder() const {
    if (_action_index == 0) return _shoulder_action_queue[0];
    return _shoulder_action_queue[_action_index - 1];
    }
    int getTargetElbow() const {
    if (_action_index == 0) return _elbow_action_queue[0];
    return _elbow_action_queue[_action_index - 1];
    }
    int getTargetGripper() const {
    if (_action_index == 0) return _gripper_action_queue[0];
    return _gripper_action_queue[_action_index - 1];
    }

    ArmController(uint8_t pin_shoulder, uint8_t pin_elbow, uint8_t pin_gripper)
    {
        _pin_shoulder = pin_shoulder;
        _pin_elbow = pin_elbow;
        _pin_gripper = pin_gripper;
        _pwm_resolution_pulse = pow(2, PWM_RESOLUTION);
    }
    void begin()
    {
        ledcSetup(SHOULDER_CHANNEL, PWM_FREQUENCY, PWM_RESOLUTION);
        ledcSetup(ELBOW_CHANNEL, PWM_FREQUENCY, PWM_RESOLUTION);
        ledcSetup(GRIPPER_CHANNEL, PWM_FREQUENCY, PWM_RESOLUTION);
        ledcAttachPin(_pin_shoulder, SHOULDER_CHANNEL);
        ledcAttachPin(_pin_elbow, ELBOW_CHANNEL);
        ledcAttachPin(_pin_gripper, GRIPPER_CHANNEL);
    }
    void arbitrary(int shoulder_angles[], int elbow_angles[], int gripper_angles[], int timeouts[], int action_num)
    {
        unsigned long now = millis();
        unsigned long cumulative_timeout = 0;
        _action_count = action_num > MAX_ACTIONS ? MAX_ACTIONS : action_num;
        for (uint8_t i = 0; i < _action_count; i++)
        {
            _shoulder_action_queue[i] = shoulder_angles[i];
            _elbow_action_queue[i] = elbow_angles[i];
            _gripper_action_queue[i] = gripper_angles[i];
            cumulative_timeout += timeouts[i];
            _action_timeout_queue[i] = now + cumulative_timeout;
        }
        _action_index = 0;
    }
    void neutral()
    {
        unsigned long now = millis();
        int shoulder_angles[1] = {0};
        int elbow_angles[1]    = {0};
        int gripper_angles[1]  = {0};
        int timeouts[1]        = {100};
        arbitrary(shoulder_angles, elbow_angles, gripper_angles, timeouts, 1);
    }
    bool action()
    {
        if (_action_index >= _action_count)
        {
            return true;
        }
        if (millis() >= _action_timeout_queue[_action_index])
        {
            ledcWrite(SHOULDER_CHANNEL, _angle2pulse(_shoulder_action_queue[_action_index]));
            ledcWrite(ELBOW_CHANNEL, _angle2pulse(_elbow_action_queue[_action_index]));
            ledcWrite(GRIPPER_CHANNEL, _angle2pulse(_gripper_action_queue[_action_index]));
            _action_index++;
        }
        return false;
    }
    int getCurrentShoulder() {
        if (_action_index < _action_count) return _shoulder_action_queue[_action_index];
        return _shoulder_action_queue[_action_count > 0 ? _action_count - 1 : 0];
    }
    int getCurrentElbow() {
        if (_action_index < _action_count) return _elbow_action_queue[_action_index];
        return _elbow_action_queue[_action_count > 0 ? _action_count - 1 : 0];
    }
    int getCurrentGripper() {
        if (_action_index < _action_count) return _gripper_action_queue[_action_index];
        return _gripper_action_queue[_action_count > 0 ? _action_count - 1 : 0];
    }

private:
    uint8_t _pin_shoulder = 19;
    uint8_t _pin_elbow = 5;
    uint8_t _pin_gripper = 22;
    uint32_t _pwm_resolution_pulse = 256;
    int32_t _shoulder_action_queue[MAX_ACTIONS] = {0};
    int32_t _elbow_action_queue[MAX_ACTIONS] = {0};
    int32_t _gripper_action_queue[MAX_ACTIONS] = {0};
    uint32_t _action_timeout_queue[MAX_ACTIONS] = {0};
    uint8_t _action_index = 0;
    uint8_t _action_count = 0;

    uint32_t _angle2pulse(int angle)
    {
        return (uint32_t)((0.0105555555 * (angle + 90) + 0.5) / 20.0 * _pwm_resolution_pulse);
    }
};

import serial
import time
import traceback
import subprocess
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("-d", "--debug", action="store_true", help="Whether to enable debug print.")
parser.add_argument("-s", "--streaming", action="store_true", help="Enable live-streaming.")

args = parser.parse_args()

ser = serial.Serial("/dev/ttyS1", 115200, timeout=0.1)

if args.debug:
    print("Debug Print enabled.")
if args.streaming:
    print("Streaming enabled.")

try:
    with subprocess.Popen(["./camera", "0", "--on-demand", "--debug" if args.debug else "", "--streaming" if args.streaming else ""], stdout=subprocess.PIPE, stderr=subprocess.PIPE) as proc:
        while True:

            # waiting for serial command
            while True:
                try:
                    data = ser.readline().decode("utf-8")
                    if not data.startswith("START,"):
                        continue
                    print(f"serial>> {data}")
                    on_demand = f"{data.split(',')[1]}\n"
                    fo = open("/home/m5stack/on_demand.txt", "w")
                    fo.write(on_demand)
                    fo.close()
                    break
                except Exception:
                    traceback.print_exc()

            # Read lines from the camera process
            while True:
                line = proc.stdout.readline().decode("utf-8").replace(" ", "").encode("utf-8")
                print(f"camera>> {line.decode('utf-8')}")
                # send to serial
                ser.write(line + b"\n")
                ser.flush()
                time.sleep(0.1)
                if b"END" in line:
                    break


except Exception:
    traceback.print_exc()

ser.close()

#include <unistd.h>
#include "opencv2/opencv.hpp"
#include "iostream"
#include "opencv2/videoio/videoio_c.h"
#include "nadjieb/mjpeg_streamer.hpp"

#include <iostream>
#include <fstream>
#include <cstdio>

#include <time.h>
#include "../tensorflow-lite/tensorflow/lite/interpreter.h"
#include "../tensorflow-lite/tensorflow/lite/kernels/register.h"
#include "../tensorflow-lite/tensorflow/lite/model.h"
#include "get_result.h"

#define CLASSIFIER_INPUT_WIDTH 320 // input image size
#define CLASSIFIER_INPUT_HEIGHT 320 // input image size
#define INFERENCE_INTERVAL_MS 200 // inference interval
#define CONF_THRESH 0.5f // confidence threshold

using MJPEGStreamer = nadjieb::MJPEGStreamer;
static bool use_debug = false;
static bool use_streaming = false;
static bool on_demand = false;
static const char* CLASS_NAMES[] = { "driver", "nipper", "needle-nose plier", "wire stripper" };
static const char* MODEL_PATH = "yolov8n.tflite"; // model file name

/*
 * Returns the current time in milliseconds
 */
long long get_current_time() {
    struct timespec ts;
    clock_gettime(CLOCK_REALTIME, &ts);
    return ts.tv_sec * 1000LL + ts.tv_nsec / 1000000;
}

/*
 * set image to input tensor
 */
void preprocess(const cv::Mat& src, float* dst) {
    cv::Mat resized, rgb, float_img;
    cv::resize(src, resized, cv::Size(CLASSIFIER_INPUT_WIDTH, CLASSIFIER_INPUT_HEIGHT));
    cv::cvtColor(resized, rgb, cv::COLOR_BGR2RGB);
    rgb.convertTo(float_img, CV_32FC3, 1.0 / 255.0);
    std::memcpy(dst, float_img.data, sizeof(float) * CLASSIFIER_INPUT_WIDTH * CLASSIFIER_INPUT_HEIGHT * 3);
}

int main(int argc, char** argv) {

    printf("OPENCV VERSION: %d_%d\n", CV_MAJOR_VERSION, CV_MINOR_VERSION);

    if (argc < 2) {
        printf("Requires one parameter (ID of the webcam).\n");
        printf("You can find these via `v4l2-ctl --list-devices`.\n");
        printf("E.g. for:\n");
        printf("    C922 Pro Stream Webcam (usb-70090000.xusb-2.1):\n");
	    printf("            /dev/video0\n");
        printf("The ID of the webcam is 0\n");
        exit(1);
    }

    for (int ix = 2; ix < argc; ix++) {
        if (strcmp(argv[ix], "--debug") == 0) {
            printf("Enabling debug mode\n");
            use_debug = true;
        }
        if (strcmp(argv[ix], "--streaming") == 0) {
            printf("Enabling streaming mode\n");
            use_streaming = true;
        }
        if (strcmp(argv[ix], "--on-demand") == 0) {
            printf("Enabling on-demand mode\n");
            on_demand = true;
        }
    }

    MJPEGStreamer streamer;
    streamer.start(80);
    std::vector<int> params = {cv::IMWRITE_JPEG_QUALITY, 90};

    // open the camera...
    cv::VideoCapture camera(atoi(argv[1]));
    if (!camera.isOpened()) {
        std::cerr << "ERROR: Could not open camera" << std::endl;
        return 1;
    } 

    std::cout << "Resolution: " <<  camera.get(CV_CAP_PROP_FRAME_WIDTH) 
	    << "x" << camera.get(CV_CAP_PROP_FRAME_HEIGHT) << std::endl;

    // tensorflow ---------------------------------------------
    // read classifier model
    auto model = tflite::FlatBufferModel::BuildFromFile(MODEL_PATH);
    if (!model) {
        std::cerr << "ERROR: failed to load model." << std::endl;
        return -1;
    }

    // create interpreter
    tflite::ops::builtin::BuiltinOpResolver resolver;
    std::unique_ptr<tflite::Interpreter> interpreter;
    tflite::InterpreterBuilder(*model, resolver)(&interpreter);

    // allocate memory
    interpreter->AllocateTensors();

    // get input tensor information
    int input = interpreter->inputs()[0];
    TfLiteIntArray* dims = interpreter->tensor(input)->dims;
    int height = dims->data[1]; // CLASSIFIER_INPUT_HEIGHT
    int width = dims->data[2]; // CLASSIFIER_INPUT_WIDTH
    int channels = dims->data[3];
    std::cout << "Input tensor shape: " << height << "x" << width << "x" << channels << std::endl;

    // ------------------------------------------------------------

    // this will contain the image from the webcam
    cv::Mat frame;
    int64_t next_frame = (int64_t)(get_current_time());
    float score_thresh = CONF_THRESH;
    std::vector<Detection> detections;
    int64_t t0;

    // display the frame until you press a key
    while (1) {
        // 1000ms. between inference
        next_frame = (int64_t)((int64_t)(get_current_time()) + INFERENCE_INTERVAL_MS);

        // capture the next frame from the webcam
        camera >> frame;

        if (on_demand) {
            // 
            const char*filename = "on_demand.txt";
            std::ifstream file(filename);
            if (file.is_open()) {
                std::string line;
                std::getline(file, line);
                score_thresh = std::stof(line);
                file.close();
                std::remove(filename);
                // printf("on_demand.txt exists: score_thresh=%f\n", score_thresh);
            } else {
                goto jump_publish;
            }
        }

        { // block

            // set image to input tensor
            float* input_data = interpreter->typed_tensor<float>(input);
            preprocess(frame, input_data);

            // inference
            t0 = (int64_t)(get_current_time());
            interpreter->Invoke();
            // printf("Inferenence time %lldms\n", (int64_t)(get_current_time()) - t0);

            // get output tensor
            int output_index = interpreter->outputs()[0];
            TfLiteTensor* output_tensor = interpreter->tensor(output_index);
            const float* output = output_tensor->data.f;
            int num_classes = output_tensor->dims->data[1] - 4; //  = data[1] - BB4
            int num_proposals = output_tensor->dims->data[2]; // BB: 
            // printf("num_classes: %d, num_proposals: %d\n", num_classes, num_proposals);

            // 
            detections = get_result(
                frame, 
                width, 
                height, 
                output, 
                num_classes, 
                num_proposals,
                score_thresh,
                use_debug
            );

            // 
            printf("RESP,%lld,%d,%f\n", t0, detections.size(), score_thresh);
            for (int k = 0; k < detections.size(); k++) {
                Detection detection = detections[k];
                cv::Rect bb = detection.box;
                printf("%s,%f,%d,%d,%d,%d\n", 
                    CLASS_NAMES[detection.class_id], 
                    detection.confidence, 
                    bb.x, 
                    bb.y, 
                    bb.width, 
                    bb.height
                );
            }
            printf("END\n");
            fflush(stdout);
        }

jump_publish:

        // draw bounding box on frame
        for (int k = 0; k < detections.size(); k++) {
            Detection detection = detections[k];
            cv::Rect bb = detection.box;
            cv::rectangle(frame, cv::Point(bb.x, bb.y), cv::Point(bb.x + bb.width, bb.y + bb.height), cv::Scalar(255, 0, 255), 2);
            cv::putText(frame, CLASS_NAMES[detection.class_id], cv::Point(bb.x, bb.y-5), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 0, 255), 1);
            cv::putText(frame, std::to_string(detection.confidence), cv::Point(bb.x, bb.y-15), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 0, 255), 1);
        }
        
        // show the image on the window
        if (use_streaming) {
            if (streamer.isRunning()) {
                std::vector<uchar> buff_bgr;
                cv::imencode(".jpg", frame, buff_bgr, params);
                streamer.publish("/stream", std::string(buff_bgr.begin(), buff_bgr.end()));
            } else {
                printf("streamer is not alive!\n");
            }
        }

    }

    streamer.stop();

    return 0;
}

#include "opencv2/opencv.hpp"

const float NMS_THRESH = 0.5f;

struct Detection {
    cv::Rect box;       // (x, y, width, height)
    float confidence;
    int class_id;
};

// IOUIntersection over Union
float iou(const cv::Rect2f& a, const cv::Rect2f& b) {
    float inter = (a & b).area();
    float union_ = a.area() + b.area() - inter;
    return inter / union_;
}

/*
 * m: num_classes
 * n: num_proposals
 * output tensor = [[
 *  [cx_0, cx_1, ..., cx_n-1],
 *  [cy_0, cy_1, ..., cy_n-1],
 *  [w_0, w_1, ..., w_n-1],
 *  [h_0, h_1, ..., h_n-1],
 *  [p_0_0, p_1_0, ..., p_m_0],
 *  :
 *  [p_0_m-1, p_1_m-1, ..., p_m-1_n-1]
 * ]]
*/

std::vector<Detection> get_result(
    cv::Mat frame,
    const int input_width,
    const int input_height,
    const float* output,
    int num_classes,
    int num_proposals,
    const float score_thresh,
    bool use_debug = false
) {

    // temporary vector for NMS
    std::vector<cv::Rect> boxes;
    std::vector<float> confidences;
    std::vector<int> class_ids;

    // image size to input size ratio
    float x_factor = frame.cols / (input_width * 1.0f);
    float y_factor = frame.rows / (input_height * 1.0f);
    // printf("x_factor: %f, y_factor: %f\n", x_factor, y_factor);
    
    for (int i = 0; i < num_proposals; ++i) {
        const float* data = output + (4 * num_proposals) + i;

        // find the class with the highest confidence score
        float max_score = 0.0f;
        int best_class_id = -1;
        for (int j = 0; j < num_classes; ++j) {
            if (data[j * num_proposals] > max_score) {
                max_score = data[j * num_proposals];
                best_class_id = j;
            }
        }

        if (max_score < score_thresh) continue;

        // get the bounding box (cx, cy, w, h)
        float cx = output[0 * num_proposals + i];
        float cy = output[1 * num_proposals + i];
        float w = output[2 * num_proposals + i];
        float h = output[3 * num_proposals + i];
        // printf("cx: %f, cy: %f, w: %f, h: %f\n", cx, cy, w, h);

        // (cx, cy, w, h) to (x, y, w, h)
        int left = static_cast<int>((cx - 0.5 * w) * x_factor * input_width);
        int top = static_cast<int>((cy - 0.5 * h) * y_factor * input_height);
        int width = static_cast<int>(w * x_factor * input_width);
        int height = static_cast<int>(h * y_factor * input_height);

        boxes.emplace_back(left, top, width, height);
        confidences.push_back(max_score);
        class_ids.push_back(best_class_id);

        if (use_debug) printf("class=%d, score=%f, bb=[%d, %d, %d, %d]\n", best_class_id, max_score, left, top, width, height);
    }

    // NMS(Non-Maximum Suppression) process
    // If the IOU of two BBs is greater than or equal to a threshold, remove the BB with the lowest reliability.
    bool nms[boxes.size()] = {false};
    for (int index = 0; index < boxes.size(); index++) {
        if (nms[index]) continue;
        cv::Rect& a = boxes[index];
        for (int other = index + 1; other < boxes.size(); other++) {
            cv::Rect& b = boxes[other];
            float _iou = iou(a, b);
            if (use_debug) printf("NMS %d vs %d: score=%f vs %f, iou=%f\n", index, other, confidences[index], confidences[other], _iou);
            if (_iou > NMS_THRESH) {
                if (confidences[index] > confidences[other]) {
                    nms[other] = true;
                    if (use_debug) printf("NMS: %d removed\n", other);
                } else {
                    nms[index] = true;
                    if (use_debug) printf("NMS: %d removed\n", index);
                    break;
                }
            }
        }
    }

    std::vector<Detection> detections;

    for (int index = 0; index < boxes.size(); index++) {
        if (nms[index]) continue;
        Detection det;
        det.box = boxes[index];
        det.confidence = confidences[index];
        det.class_id = class_ids[index];
        detections.push_back(det);
    }

    return detections;

}