Refactor streamwriter, require explictly specifying the streams to wr…

…ite out
wusopp · Jan 23, 2021 · ef98c3c · ef98c3c
1 parent e2a7e01
commit ef98c3c
Show file tree

Hide file tree

Showing 7 changed files with 215 additions and 142 deletions.
diff --git a/camera.cpp b/camera.cpp
@@ -22,7 +22,6 @@
 
 #include "camera.hpp"
 #include "util.hpp"
-#include "stream_writer.hpp"
 #include "image_util.hpp"
 
 Camera::Camera(int camera, Database& db) : id(camera), db(db), stream(cfg), fm(db, cfg) {
@@ -32,6 +31,16 @@ Camera::Camera(int camera, Database& db) : id(camera), db(db), stream(cfg), fm(d
     alive = true;
     watchdog = false;
     startup = true;
+
+    //variables for cutting files...
+    last_cut = av_make_q(0, 1);
+    int seconds_per_file_raw = cfg.get_value_int("seconds-per-file");
+    if (seconds_per_file_raw <= 0){
+        LWARN("seconds-per-file was invalid");
+        seconds_per_file_raw = DEFAULT_SECONDS_PER_FILE;
+    }
+    seconds_per_file = av_make_q(seconds_per_file_raw, 1);
+
 }
 Camera::~Camera(){
 
@@ -52,89 +61,58 @@ void Camera::run(void){
     startup = false;
 
     LINFO("Camera " + std::to_string(id) + " connected...");
-    long int file_id;
-    std::string new_output = fm.get_next_path(file_id, id, stream.get_start_time());
+    std::string out_filename = fm.get_next_path(file_id, id, stream.get_start_time());
 
-    StreamWriter out = StreamWriter(cfg, new_output, stream);
+    StreamWriter out = StreamWriter(cfg);
+    out.change_path(out_filename);
+    out.copy_streams(stream);
     out.open();
 
     AVPacket pkt;
     bool valid_keyframe = false;
 
+    bool decode = true;//should we decode frames?
     //allocate a frame (we only allocate when we want to decode, an unallocated frame means we skip decoding)
     AVFrame *frame = NULL;
-    frame = av_frame_alloc();
-    if (!frame){
-        LWARN("Failed to allocate frame");
-        frame = NULL;
-    }
-
-    //variables for cutting files...
-    AVRational last_cut = av_make_q(0, 1);
-    int seconds_per_file_raw = cfg.get_value_int("seconds-per-file");
-    if (seconds_per_file_raw <= 0){
-        LWARN("seconds-per-file was invalid");
-        seconds_per_file_raw = DEFAULT_SECONDS_PER_FILE;
+    if (decode){
+        frame = av_frame_alloc();
+        if (!frame){
+            LWARN("Failed to allocate frame");
+            frame = NULL;
+        }
     }
-    AVRational seconds_per_file = av_make_q(seconds_per_file_raw, 1);
 
     //used for calculating shutdown time for the last segment
     long last_pts = 0;
     long last_stream_index = 0;
-    int frame_count = 0;
     while (!shutdown && stream.get_next_frame(pkt)){
-        frame_count++;
         watchdog = true;
         last_pts = pkt.pts;
         last_stream_index = pkt.stream_index;
 
-        if (pkt.flags & AV_PKT_FLAG_KEY && stream.get_format_context()->streams[pkt.stream_index]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO){
-            //calculate the seconds:
-            AVRational sec = av_mul_q(av_make_q(pkt.dts, 1), stream.get_format_context()->streams[pkt.stream_index]->time_base);//current time..
-            sec = av_sub_q(sec, last_cut);
-            if (av_cmp_q(sec, seconds_per_file) == 1 || sec.num < 0){
-                //cutting the video
-                struct timeval start;
-                Util::compute_timestamp(stream.get_start_time(), start, pkt.pts, stream.get_format_context()->streams[pkt.stream_index]->time_base);
-                out.close();
-                fm.update_file_metadata(file_id, start);
-                if (sec.num < 0){
-                    //this will cause a discontinuity which will be picked up and cause it to play correctly
-                    start.tv_usec += 1000;
-                }
-                new_output = fm.get_next_path(file_id, id, start);
-                out.change_path(new_output);
-                if (!out.open()){
-                    shutdown = true;
-                }
-                //save out this position
-                last_cut = av_mul_q(av_make_q(pkt.dts, 1), stream.get_format_context()->streams[pkt.stream_index]->time_base);
-            }
-        }        
-        if (valid_keyframe || (pkt.flags & AV_PKT_FLAG_KEY && stream.get_format_context()->streams[pkt.stream_index]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)){
-            out.write(pkt);//log it
-            valid_keyframe = true;
-            //LINFO("Got Frame " + std::to_string(id));
-        }
-
         //decode the packets
         if (frame){
-            if (stream.decode_packet(pkt)){
-                while (stream.get_decoded_frame(pkt.stream_index, frame)){
-                    if (stream.get_format_context()->streams[pkt.stream_index]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO){
-                        if (frame_count % 30 == 0){
-                            ImageUtil iu(db, cfg);
-                            struct timeval framestart;
-                            Util::compute_timestamp(stream.get_start_time(), framestart, pkt.pts, stream.get_format_context()->streams[pkt.stream_index]->time_base);
-                            std::string filename = "test";
-                            LWARN("Writing image");
-                            iu.write_frame_jpg(frame, filename, &framestart);
-                        }
+            if (stream.is_video(pkt)){
+                if (stream.decode_packet(pkt)){
+                    while (stream.get_decoded_frame(pkt.stream_index, frame)){
                         LWARN("Decoded Video Frame");
                     }
                 }
+            } else if (stream.is_audio(pkt)){
+                if (stream.decode_packet(pkt)){
+                    while (stream.get_decoded_frame(pkt.stream_index, frame)){
+                        LWARN("Decoded Audio Frame");
+                    }
+                }
             }
         }
+
+        cut_video(pkt, out);
+        if (valid_keyframe || (pkt.flags & AV_PKT_FLAG_KEY && stream.is_video(pkt))){
+            out.write(pkt, stream.get_stream(pkt));//log it
+            valid_keyframe = true;
+            //LINFO("Got Frame " + std::to_string(id));
+        }
 
         stream.unref_frame(pkt);
     }
@@ -176,3 +154,31 @@ void Camera::set_thread_id(std::thread::id tid){
 std::thread::id Camera::get_thread_id(void){
     return thread_id;
 }
+
+void Camera::cut_video(AVPacket &pkt, StreamWriter &out){
+    if (pkt.flags & AV_PKT_FLAG_KEY && stream.is_video(pkt)){
+        //calculate the seconds:
+        AVRational sec = av_mul_q(av_make_q(pkt.dts, 1), stream.get_format_context()->streams[pkt.stream_index]->time_base);//current time..
+        sec = av_sub_q(sec, last_cut);
+        if (av_cmp_q(sec, seconds_per_file) == 1 || sec.num < 0){
+            //cutting the video
+            struct timeval start;
+            stream.timestamp(pkt, start);
+            out.close();
+            fm.update_file_metadata(file_id, start);
+            if (sec.num < 0){
+                //this will cause a discontinuity which will be picked up and cause it to play correctly
+                start.tv_usec += 1000;
+            }
+            std::string out_filename = fm.get_next_path(file_id, id, start);
+            out.change_path(out_filename);
+            out.copy_streams(stream);
+            if (!out.open()){
+                shutdown = true;
+            }
+            //save out this position
+            last_cut = av_mul_q(av_make_q(pkt.dts, 1), stream.get_format_context()->streams[pkt.stream_index]->time_base);
+        }
+    }
+
+}
diff --git a/camera.hpp b/camera.hpp
@@ -25,6 +25,7 @@
 #include "database.hpp"
 #include "chiton_config.hpp"
 #include "stream_unwrap.hpp"
+#include "stream_writer.hpp"
 #include "file_manager.hpp"
 #include <atomic>
 #include <thread>
@@ -58,5 +59,12 @@ class Camera {
     std::atomic_bool shutdown;//signals that we should exit
     std::atomic_bool startup;//used to identify if we are in an extended wait due to startup
     std::thread::id thread_id;//used for tracking our thread
+
+    AVRational last_cut;
+    AVRational seconds_per_file;
+
+    long int file_id;//database id of current file we are writing to
+    //check if packet is a keyframe and switch the filename as needed
+    void cut_video(AVPacket &pkt, StreamWriter &out);
 };
 #endif
diff --git a/export.cpp b/export.cpp
@@ -108,10 +108,13 @@ void Export::run_job(void){
     }
 
     StreamUnwrap in = StreamUnwrap(camera_cfg);
-    StreamWriter out = StreamWriter(camera_cfg, path, in);
+    StreamWriter out = StreamWriter(camera_cfg);
+    out.change_path(path);
+
     long total_time_target = endtime - starttime;
 
     long reserved_bytes = 0;
+    bool output_opened = false;
     while (res->next_row()){
         std::string segment = fm.get_path(res->get_field_long(0), res->get_field(1));
         LDEBUG("Exporting " + segment);
@@ -130,12 +133,21 @@ void Export::run_job(void){
             reserved_bytes -= seg_size;
         }
 
-        AVPacket pkt;
-        in.connect();
-        out.open();
+        if (!in.connect()){
+            //try the next one...
+            continue;
+        }
+
+        if (!output_opened){
+            if (!out.copy_streams(in)){
+                continue;
+            }
+            output_opened = out.open();
+        }
 
+        AVPacket pkt;
         while (in.get_next_frame(pkt)){
-            out.write(pkt);
+            out.write(pkt, in.get_stream(pkt));
             in.unref_frame(pkt);
         }
         in.close();
@@ -145,10 +157,12 @@ void Export::run_job(void){
         }
 
         //compute the progress
-        long new_progress = 100 - (100*(endtime - res->get_field_long(2))) / total_time_target;
-        if (new_progress > progress && new_progress < 100){//this can result in number over 100, we don't put 100 in until it's actually done
-            progress = new_progress;
-            update_progress();
+        if (res->get_field(2) != "NULL"){//do not compute the progress if we don't know how long this is
+            long new_progress = 100 - (100*(endtime - res->get_field_long(2))) / total_time_target;
+            if (new_progress > progress && new_progress < 100){//this can result in number over 100, we don't put 100 in until it's actually done
+                progress = new_progress;
+                update_progress();
+            }
         }
     }
     out.close();

diff --git a/stream_unwrap.cpp b/stream_unwrap.cpp
@@ -350,3 +350,23 @@ bool StreamUnwrap::get_decoded_frame(int stream, AVFrame *frame){
         return false;
     }
 }
+
+void StreamUnwrap::timestamp(const AVPacket &packet, struct timeval &time){
+    Util::compute_timestamp(get_start_time(), time, packet.pts, input_format_context->streams[pkt.stream_index]->time_base);
+}
+
+bool StreamUnwrap::is_audio(const AVPacket &packet){
+    return input_format_context->streams[pkt.stream_index]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO;
+}
+
+bool StreamUnwrap::is_video(const AVPacket &packet){
+    return input_format_context->streams[pkt.stream_index]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO;
+}
+
+AVStream *StreamUnwrap::get_stream(const AVPacket &packet){
+    return get_stream(packet.stream_index);
+}
+
+AVStream *StreamUnwrap::get_stream(const int id){
+    return input_format_context->streams[id];
+}
diff --git a/stream_unwrap.hpp b/stream_unwrap.hpp
@@ -47,6 +47,11 @@ class StreamUnwrap {
     bool decode_packet(AVPacket &packet);//reads packet and decodes it
     bool get_decoded_frame(int stream, AVFrame *frame);//gets the next decoded frame
     void unref_frame(AVPacket &packet);//free resources from frame
+    void timestamp(const AVPacket &packet, struct timeval &time);//write the actual timestamp of packet to time
+    bool is_audio(const AVPacket &packet);//return true if packet is audio packet
+    bool is_video(const AVPacket &packet);//return true if packet is video packet
+    AVStream *get_stream(const AVPacket &packet);//return a pointer to the stream that packet is part of
+    AVStream *get_stream(const int id);//return a pointer to the stream given the stream index
     const struct timeval& get_start_time(void);
 
 private: