642 lines
18 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package main
import (
"bytes"
"errors"
"log"
"os"
"strings"
"sync"
"time"
"github.com/Eyevinn/mp4ff/mp4"
"github.com/bluenviron/gortsplib/v4"
"github.com/bluenviron/gortsplib/v4/pkg/base"
"github.com/bluenviron/gortsplib/v4/pkg/description"
"github.com/bluenviron/gortsplib/v4/pkg/format"
"github.com/bluenviron/gortsplib/v4/pkg/format/rtph264"
"github.com/gen2brain/aac-go"
"github.com/pion/rtp"
"github.com/zaf/g711"
)
var (
currentVideoFile *os.File
currentAudioFile *os.File
currentMP4File *os.File
wg sync.WaitGroup
fileMu sync.Mutex
formaH264Mu sync.Mutex
formaH264 *format.H264
currentSegmentStarted bool // Флаг указывает записан ли уже IDR (с SPS/PPS) в текущий сегмент
)
func main() {
// Парсинг URL RTSP
u, err := base.ParseURL("rtsp://intercom-video-1.insit.ru/dp-gfahybswjoendkcpbaqvgkeizruudbhsfdr")
if err != nil {
log.Fatalln("Ошибка парсинга URL:", err)
}
// Инициализация клиента и подключение к камере
client := gortsplib.Client{}
err = client.Start(u.Scheme, u.Host)
if err != nil {
log.Fatalln("Ошибка соединения:", err)
}
defer client.Close()
// Получение описания
desc, _, err := client.Describe(u)
if err != nil {
log.Fatalln("Ошибка получения описания:", err)
}
// Установка настроек
err = client.SetupAll(desc.BaseURL, desc.Medias)
if err != nil {
log.Fatalln("Ошибка настройки:", err)
}
// Проверка форматов для видео- и аудио-потоков
var h264Format *format.H264
var g711Format *format.G711
for _, media := range desc.Medias {
switch media.Formats[0].(type) {
case *format.H264:
h264Format = media.Formats[0].(*format.H264)
case *format.G711:
g711Format = media.Formats[0].(*format.G711)
}
}
if h264Format == nil || g711Format == nil {
log.Fatalln("Форматы не найдены")
}
// Создание декодера для видео-потока H264
decoderH264, err := h264Format.CreateDecoder()
if err != nil {
log.Fatalln("Ошибка создания декодера H264:", err)
}
// Обработка RTP-пакетов: для видео сбор NAL-ов и запись в файл с префиксом, для аудио запись PCM-данных.
client.OnPacketRTPAny(func(medi *description.Media, forma format.Format, pkt *rtp.Packet) {
switch f := forma.(type) {
case *format.H264:
formaH264Mu.Lock()
formaH264 = f
formaH264Mu.Unlock()
nalus, err := decoderH264.Decode(pkt)
if err != nil && !errors.Is(err, rtph264.ErrMorePacketsNeeded) {
log.Printf("Ошибка декодирования H264: %v", err)
return
}
fileMu.Lock()
if currentVideoFile != nil {
// Если сегмент ещё не начат, ожидается появление ключевого кадра (IDR).
if !currentSegmentStarted {
var isIDR bool
for _, nalu := range nalus {
if len(nalu) > 0 && (nalu[0]&0x1F) == 5 {
isIDR = true
break
}
}
if !isIDR {
// Пакет не записывается, если в нём нет ключевого кадра.
fileMu.Unlock()
return
}
// При получении IDR происходит вставка SPS и PPS (с префиксами) в начало сегмента.
if len(f.SPS) > 0 {
if _, err := currentVideoFile.Write([]byte{0x00, 0x00, 0x00, 0x01}); err != nil {
log.Printf("Ошибка записи стартового кода SPS: %v", err)
}
if _, err := currentVideoFile.Write(f.SPS); err != nil {
log.Printf("Ошибка записи SPS: %v", err)
}
}
if len(f.PPS) > 0 {
if _, err := currentVideoFile.Write([]byte{0x00, 0x00, 0x00, 0x01}); err != nil {
log.Printf("Ошибка записи стартового кода PPS: %v", err)
}
if _, err := currentVideoFile.Write(f.PPS); err != nil {
log.Printf("Ошибка записи PPS: %v", err)
}
}
currentSegmentStarted = true
}
// Запись каждой NAL-единицы с префиксом
for _, nalu := range nalus {
if _, err := currentVideoFile.Write([]byte{0x00, 0x00, 0x00, 0x01}); err != nil {
log.Printf("Ошибка записи стартового кода NALU: %v", err)
}
if _, err := currentVideoFile.Write(nalu); err != nil {
log.Printf("Ошибка записи NALU: %v", err)
}
}
}
fileMu.Unlock()
case *format.G711:
if strings.Contains(forma.RTPMap(), "PCMA") {
sampleG711a := g711.DecodeAlaw(pkt.Payload)
fileMu.Lock()
if currentAudioFile != nil {
if _, err := currentAudioFile.Write(sampleG711a); err != nil {
log.Printf("Ошибка записи аудио: %v", err)
}
}
fileMu.Unlock()
} else if strings.Contains(forma.RTPMap(), "PCMU") {
sampleG711u := g711.DecodeUlaw(pkt.Payload)
fileMu.Lock()
if currentAudioFile != nil {
if _, err := currentAudioFile.Write(sampleG711u); err != nil {
log.Printf("Ошибка записи аудио: %v", err)
}
}
fileMu.Unlock()
} else {
log.Println("Аудиокодек не идентифицирован")
}
}
})
// Запуск воспроизведения
_, err = client.Play(nil)
if err != nil {
log.Fatalln("Ошибка запуска воспроизведения:", err)
}
// Параметр записи
period := time.Hour
// Ожидание начала следующего часа
now := time.Now()
nextSegment := now.Truncate(time.Hour).Add(time.Hour)
waitDuration := nextSegment.Sub(now)
log.Printf("Ожидание до начала записи: %v\n", waitDuration)
time.Sleep(waitDuration)
log.Println("Начало записи фрагмента")
// Создание начальных файлов
initialTimestamp := time.Now().Format("15-04_02-01-2006")
videoFilename := initialTimestamp + ".h264"
audioFilename := initialTimestamp + ".pcm"
mp4Filename := initialTimestamp + ".mp4"
fileVideo, err := os.Create(videoFilename)
if err != nil {
log.Fatalln("Ошибка создания видеофайла:", err)
}
fileAudio, err := os.Create(audioFilename)
if err != nil {
log.Fatalln("Ошибка создания аудиофайла:", err)
}
fileMP4, err := os.Create(mp4Filename)
if err != nil {
log.Fatalln("Ошибка создания mp4 файла:", err)
}
fileMu.Lock()
currentVideoFile = fileVideo
currentAudioFile = fileAudio
currentMP4File = fileMP4
// Начало нового сегмента флаг сбрасывается и ожидается первый IDR
currentSegmentStarted = false
fileMu.Unlock()
// Тикер для периодической ротации файлов
ticker := time.NewTicker(period)
defer ticker.Stop()
// Горутина для смены файлов
go func() {
for range ticker.C {
newTimestamp := time.Now().Format("15-04_02-01-2006")
newVideoFilename := newTimestamp + ".h264"
newAudioFilename := newTimestamp + ".pcm"
newMP4Filename := newTimestamp + ".mp4"
newVideoFile, err := os.Create(newVideoFilename)
if err != nil {
log.Printf("Ошибка создания видеофайла для фрагмента %s: %v", newTimestamp, err)
continue
}
newAudioFile, err := os.Create(newAudioFilename)
if err != nil {
log.Printf("Ошибка создания аудиофайла для фрагмента %s: %v", newTimestamp, err)
continue
}
newMP4File, err := os.Create(newMP4Filename)
if err != nil {
log.Printf("Ошибка создания mp4 файла для фрагмента %s: %v", newTimestamp, err)
}
// Переключение на новые файлы
fileMu.Lock()
oldVideoFile := currentVideoFile
oldAudioFile := currentAudioFile
oldMP4File := currentMP4File
oldVideoFilename := videoFilename
oldAudioFilename := audioFilename
oldMP4Filename := mp4Filename
currentVideoFile = newVideoFile
currentAudioFile = newAudioFile
currentMP4File = newMP4File
// Новый сегмент ещё не начат флаг сбрасывается и ожидается первый IDR
currentSegmentStarted = false
// Обновление имен для следующей итерации
videoFilename = newVideoFilename
audioFilename = newAudioFilename
mp4Filename = newMP4Filename
fileMu.Unlock()
// Закрытие старых файлов
oldVideoFile.Close()
oldAudioFile.Close()
oldMP4File.Close()
log.Println("Созданы новые файлы для записи:", videoFilename, audioFilename, mp4Filename)
// Горутина для объединения видео и аудио в контейнер MP4
go func(oldVideoFilename, oldAudioFilename, oldMP4Filename string) {
formaH264Mu.Lock()
h264Params := formaH264
formaH264Mu.Unlock()
// Конвертация PCM в AAC
aacFilename := strings.Replace(oldAudioFilename, ".pcm", ".aac", 1)
if err := convertPCMtoAAC(oldAudioFilename, aacFilename); err != nil {
log.Printf("Ошибка конвертации аудио: %v", err)
return
}
log.Printf("Файл %s успешно конвертирован в %s\n", oldAudioFilename, aacFilename)
// Создание MP4 файл
if err := createMP4File(oldVideoFilename, aacFilename, oldMP4Filename, h264Params); err != nil {
log.Printf("Ошибка создания MP4: %v", err)
} else {
log.Printf("Файлы %s, %s успешно объединены в контейнер MP4: %s\n",
oldVideoFilename, aacFilename, oldMP4Filename)
}
// Удаление временных файлов
os.Remove(oldVideoFilename)
os.Remove(aacFilename)
os.Remove(oldAudioFilename)
log.Printf("Файлы %s, %s, %s успешно удалены\n", oldVideoFilename, aacFilename, oldAudioFilename)
}(oldVideoFilename, oldAudioFilename, oldMP4Filename)
}
}()
wg.Add(1)
wg.Wait()
}
// Sample единый тип для описания сэмпла (для видео и аудио)
type Sample struct {
Data []byte
Offset uint64
Size uint32
Dur uint32
IsSync bool
}
// convertPCMtoAAC конвертирует PCM-файл в AAC-файл.
func convertPCMtoAAC(aFilename, outputFile string) error {
if outputFile == "" {
outputFile = strings.Replace(aFilename, ".pcm", ".aac", 1)
}
audioFileAAC, err := os.Create(outputFile)
if err != nil {
log.Println("Ошибка создания файла AAC:", err)
return err
}
defer audioFileAAC.Close()
audioFilePCM, err := os.Open(aFilename)
if err != nil {
log.Println("Ошибка открытия PCM файла:", err)
return err
}
defer audioFilePCM.Close()
opts := &aac.Options{
SampleRate: 8000,
NumChannels: 1,
}
encoderAAC, err := aac.NewEncoder(audioFileAAC, opts)
if err != nil {
log.Println("Ошибка создания инкодера AAC:", err)
return err
}
defer encoderAAC.Close()
if err = encoderAAC.Encode(audioFilePCM); err != nil {
log.Println("Ошибка инкодирования в AAC:", err)
return err
}
if err = os.Remove(aFilename); err != nil {
log.Println("Ошибка удаления PCM файла:", err)
} else {
log.Printf("PCM файл %s успешно удалён.", aFilename)
}
return nil
}
// createMP4File упаковывает видео (H.264) и аудио (AAC) файлы в MP4.
func createMP4File(videoFile, audioFile, outputFile string, formaH264 *format.H264) error {
videoData, err := os.ReadFile(videoFile)
if err != nil {
return err
}
videoSamples := parseH264Samples(videoData)
audioData, err := os.ReadFile(audioFile)
if err != nil {
return err
}
audioSamples := parseAACSamples(audioData)
f := mp4.NewFile()
videoTrack := createVideoTrack(formaH264.SPS, formaH264.PPS)
audioTrack := createAudioTrack(8000, 1)
addSamplesToVideoTrack(videoTrack, videoSamples)
addSamplesToAudioTrack(audioTrack, audioSamples, 8000)
moov := mp4.NewMoovBox()
moov.AddChild(videoTrack)
moov.AddChild(audioTrack)
f.AddChild(moov, 0)
mdat := mp4.MdatBox{}
mdat.Data = append(videoData, audioData...)
f.AddChild(&mdat, 0)
outFile, err := os.Create(outputFile)
if err != nil {
return err
}
defer outFile.Close()
return f.Encode(outFile)
}
// createVideoTrack создаёт видео-трек на основе SPS и PPS.
func createVideoTrack(sps, pps []byte) *mp4.TrakBox {
sps = stripAnnexB(sps)
pps = stripAnnexB(pps)
if len(sps) < 4 {
log.Fatalln("SPS слишком короткий или пустой")
}
if len(pps) < 1 {
log.Fatalln("PPS слишком короткий или пустой")
}
trak := mp4.NewTrakBox()
timescale := uint32(90000)
mdhd := &mp4.MdhdBox{
Timescale: timescale,
Language: convertLanguage("und"),
}
hdlr := &mp4.HdlrBox{
HandlerType: "vide",
Name: "VideoHandler",
}
avc1 := mp4.NewVisualSampleEntryBox("avc1")
avc1.Width = 1280
avc1.Height = 720
avcC, err := mp4.CreateAvcC([][]byte{sps}, [][]byte{pps}, true)
if err != nil {
log.Fatalf("ошибка создания avcC: %v", err)
}
avcC.AVCLevelIndication = 0x1f
avc1.AddChild(avcC)
stbl := mp4.NewStblBox()
stsd := mp4.NewStsdBox()
stsd.AddChild(avc1)
stbl.AddChild(stsd)
stbl.AddChild(&mp4.SttsBox{})
stbl.AddChild(&mp4.StscBox{})
stbl.AddChild(&mp4.StszBox{})
stbl.AddChild(&mp4.StcoBox{})
trak.Mdia = &mp4.MdiaBox{
Mdhd: mdhd,
Hdlr: hdlr,
Minf: &mp4.MinfBox{
Stbl: stbl,
},
}
return trak
}
// stripAnnexB удаляет стартовый код Annex-B из NAL-единицы.
func stripAnnexB(nalu []byte) []byte {
if len(nalu) >= 4 && bytes.Equal(nalu[:4], []byte{0x00, 0x00, 0x00, 0x01}) {
return nalu[4:]
}
if len(nalu) >= 3 && bytes.Equal(nalu[:3], []byte{0x00, 0x00, 0x01}) {
return nalu[3:]
}
return nalu
}
// createAudioTrack создаёт аудио-трек.
func createAudioTrack(sampleRate, channels int) *mp4.TrakBox {
trak := mp4.NewTrakBox()
timescale := uint32(sampleRate)
mdhd := &mp4.MdhdBox{
Timescale: timescale,
Language: convertLanguage("und"),
}
hdlr := &mp4.HdlrBox{
HandlerType: "soun",
Name: "SoundHandler",
}
mp4a := mp4.NewAudioSampleEntryBox("mp4a")
mp4a.ChannelCount = uint16(channels)
mp4a.SampleRate = uint16(sampleRate)
mp4a.SampleSize = 16
asc := getAACConfig(sampleRate, channels)
esds := createESDSBox(asc)
mp4a.AddChild(esds)
stbl := mp4.NewStblBox()
stsd := mp4.NewStsdBox()
stsd.AddChild(mp4a)
stbl.AddChild(stsd)
stbl.AddChild(&mp4.SttsBox{})
stbl.AddChild(&mp4.StscBox{})
stbl.AddChild(&mp4.StszBox{})
stbl.AddChild(&mp4.StcoBox{})
trak.Mdia = &mp4.MdiaBox{
Mdhd: mdhd,
Hdlr: hdlr,
Minf: &mp4.MinfBox{
Stbl: stbl,
},
}
return trak
}
// createESDSBox создаёт ESDS-бокс для AAC.
func createESDSBox(asc []byte) *mp4.EsdsBox {
return &mp4.EsdsBox{
ESDescriptor: mp4.ESDescriptor{
DecConfigDescriptor: &mp4.DecoderConfigDescriptor{
ObjectType: 0x40,
StreamType: 0x05,
DecSpecificInfo: &mp4.DecSpecificInfoDescriptor{
DecConfig: asc,
},
},
},
}
}
func getAACConfig(sampleRate, channels int) []byte {
// Пример конфигурации для 8000 Гц, моно
return []byte{0x12, 0x10}
}
// addSamplesToVideoTrack заполняет таблицы сэмплов видео-трека.
func addSamplesToVideoTrack(trak *mp4.TrakBox, samples []Sample) {
stbl := trak.Mdia.Minf.Stbl
stts := mp4.SttsBox{}
for _, s := range samples {
stts.SampleCount = []uint32{1}
stts.SampleTimeDelta = []uint32{s.Dur}
}
stsz := mp4.StszBox{}
for _, s := range samples {
stsz.SampleSize = []uint32{s.Size}
}
stss := mp4.StssBox{}
for i, s := range samples {
if s.IsSync {
stss.SampleNumber = []uint32{uint32(i + 1)}
}
}
stbl.AddChild(&stts)
stbl.AddChild(&stsz)
stbl.AddChild(&stss)
}
// addSamplesToAudioTrack заполняет таблицы сэмплов аудио-трека.
func addSamplesToAudioTrack(trak *mp4.TrakBox, samples []Sample, sampleRate int) {
stbl := trak.Mdia.Minf.Stbl
stts := mp4.SttsBox{}
for _, s := range samples {
stts.SampleCount = []uint32{1}
stts.SampleTimeDelta = []uint32{s.Dur}
}
stsz := mp4.StszBox{}
for _, s := range samples {
stsz.SampleSize = []uint32{uint32(len(s.Data))}
}
stbl.AddChild(&stts)
stbl.AddChild(&stsz)
}
// parseH264Samples парсит H.264 поток на сэмплы.
func parseH264Samples(data []byte) []Sample {
var samples []Sample
start := 0
for i := 0; i < len(data)-3; i++ {
if bytes.Equal(data[i:i+4], []byte{0x00, 0x00, 0x00, 0x01}) {
if start != 0 {
sampleData := data[start:i]
samples = append(samples, Sample{
Size: uint32(len(sampleData)),
Dur: 3600,
IsSync: isKeyFrame(sampleData),
})
}
start = i + 4
}
}
if start < len(data) {
sampleData := data[start:]
samples = append(samples, Sample{
Size: uint32(len(sampleData)),
Dur: 3600,
IsSync: isKeyFrame(sampleData),
})
}
return samples
}
// isKeyFrame определяет, является ли NALU ключевым (IDR).
func isKeyFrame(nalu []byte) bool {
if len(nalu) == 0 {
return false
}
nalType := nalu[0] & 0x1F
return nalType == 5
}
// parseAACSamples парсит AAC данные и возвращает сэмплы.
func parseAACSamples(data []byte) []Sample {
var samples []Sample
frameSize := 1024 // примерный размер AAC-фрейма
for i := 0; i < len(data); i += frameSize {
end := i + frameSize
if end > len(data) {
end = len(data)
}
samples = append(samples, Sample{
Data: data[i:end],
Dur: 1024,
Size: uint32(end - i),
})
}
return samples
}
// convertLanguage преобразует строку языка в 16-битное значение.
func convertLanguage(lang string) uint16 {
if len(lang) != 3 {
return 0x55C4
}
b1 := lang[0] - 0x60
b2 := lang[1] - 0x60
b3 := lang[2] - 0x60
return uint16((b1 << 10) | (b2 << 5) | b3)
}