Fixed memory leak in bitstream analyze stage

osai-ai · Oct 28, 2019 · 968b646 · 968b646
1 parent db7b55c
commit 968b646
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 3 deletions.
diff --git a/docker/Dockerfile_cu10 b/docker/Dockerfile_cu10
@@ -37,7 +37,8 @@ RUN git clone --depth 1 -b release/4.2 --single-branch https://github.com/FFmpeg
 RUN pip3 install --no-cache-dir \
     twine==1.13.0 \
     awscli==1.16.194 \
-    numpy==1.16.4
+    numpy==1.16.4 \
+    packaging
 
 # Install PyTorch
 RUN pip3 install --no-cache-dir \

diff --git a/include/Parser.h b/include/Parser.h
@@ -152,7 +152,7 @@ class Parser {
 	Bitstream filter for converting mp4->h264
 	*/
 	AVBitStreamFilterContext* bitstreamFilter;
-	std::shared_ptr<AVPacket> NALu;
+	AVPacket* NALu;
 	/*
 	Instance of Logger class
 	*/

diff --git a/src/Parser.cpp b/src/Parser.cpp
@@ -276,6 +276,10 @@ int Parser::Analyze(AVPacket* package) {
 		frameNumValue = frame_num;
 		POC = pic_order_cnt_lsb;
 	}
+
+	av_freep(&NALu->data);
+	av_packet_free_side_data(NALu);
+	av_free_packet(NALu);
 	return errorBitstream;
 }
 
@@ -310,7 +314,9 @@ int Parser::Init(ParserParameters& input, std::shared_ptr<Logger> logger) {
 		//Write file header
 		sts = avformat_write_header(dumpContext, NULL);
 	}
-	NALu = std::make_shared<AVPacket>();
+	NALu = new AVPacket();
+	av_init_packet(NALu);
+
 	bitstreamFilter = av_bitstream_filter_init("h264_mp4toannexb");
 
 	lastFrame = std::make_pair(new AVPacket(), false);
@@ -402,5 +408,9 @@ void Parser::Close() {
 	}
 	av_packet_unref(lastFrame.first);
 	delete lastFrame.first;
+
+	av_packet_unref(NALu);
+	delete NALu;
+
 	isClosed = true;
 }