Merge pull request #1 from r-arpitgupta/vtt-segmentation

vtt segmentation

Author: r-arpitgupta

astisub/main.go

@@ -2,7 +2,10 @@ package main
 
 import (
 	"flag"
+	"fmt"
 	"log"
+	"strconv"
+	"strings"
 
 	"github.com/asticode/go-astikit"
 	"github.com/asticode/go-astisub"
@@ -19,6 +22,11 @@ var (
 	teletextPage     = flag.Int("p", 0, "the teletext page")
 	outputPath       = flag.String("o", "", "the output path")
 	syncDuration     = flag.Duration("s", 0, "the sync duration")
+	segmentationType = flag.String("st", "UNIFIED", "the segmentation type UNIFIED/SPECIFIED. In unified"+
+		" segmentation all segment have same duration whereas specified have user given duration for each segment.")
+	segmentDuration  = flag.Float64("sd", 5, "segmentation duration for unified segmentation type")
+	segmentDurations = flag.String("sds", "", "segment durations for all segments seperated by comma")
+	webvttOffset     = flag.Float64("wo", 0, "webvtt offset for synchronization of segment in hls")
 )
 
 func main() {
@@ -134,6 +142,27 @@ func main() {
 		if err = sub.Write(*outputPath); err != nil {
 			log.Fatalf("%s while writing to %s", err, *outputPath)
 		}
+	case "vtt-segment":
+		// go run main.go vtt-segment -i <input_path>  -st SPECIFIED -sds "2,4,10,12,13" -o "seg-%03d.vtt" -wo 10
+		if *segmentationType == "SPECIFIED" && segmentDurations == nil {
+			log.Fatalf("segment duration must be present for specified segmentation type")
+		}
+		var sds []float64
+		if *segmentDurations != "" {
+			for _, segDuration := range strings.Split(*segmentDurations, ",") {
+				segDur, err := strconv.ParseFloat(segDuration, 64)
+				if err != nil {
+					log.Fatalf("%s unable to parse seg duration %s", err, segDuration)
+				}
+				sds = append(sds, segDur)
+			}
+		}
+		segmentedSubs := sub.Segment(*segmentationType, *segmentDuration, sds)
+		for idx, segmentedSub := range segmentedSubs {
+			if err = segmentedSub.WriteToWebVTTFile(fmt.Sprintf(*outputPath, idx), *webvttOffset); err != nil {
+				log.Fatalf("%s while writing to %s", err, *outputPath)
+			}
+		}
 	default:
 		log.Fatalf("Invalid subcommand %s", cmd)
 	}

subtitles.go

@@ -318,28 +318,28 @@ func (sa *StyleAttributes) propagateTTMLAttributes() {
 		sa.WebVTTAlign = *sa.TTMLTextAlign
 	}
 	if sa.TTMLExtent != nil {
-		//region settings
-		lineHeight := 5 //assuming height of line as 5.33vh
+		// region settings
+		lineHeight := 5 // assuming height of line as 5.33vh
 		dimensions := strings.Split(*sa.TTMLExtent, " ")
 		if len(dimensions) > 1 {
 			sa.WebVTTWidth = dimensions[0]
 			if height, err := strconv.Atoi(strings.ReplaceAll(dimensions[1], "%", "")); err == nil {
 				sa.WebVTTLines = height / lineHeight
 			}
-			//cue settings
-			//default TTML WritingMode is lrtb i.e. left to right, top to bottom
+			// cue settings
+			// default TTML WritingMode is lrtb i.e. left to right, top to bottom
 			sa.WebVTTSize = dimensions[1]
 			if sa.TTMLWritingMode != nil && strings.HasPrefix(*sa.TTMLWritingMode, "tb") {
 				sa.WebVTTSize = dimensions[0]
 			}
 		}
 	}
 	if sa.TTMLOrigin != nil {
-		//region settings
+		// region settings
 		sa.WebVTTRegionAnchor = "0%,0%"
 		sa.WebVTTViewportAnchor = strings.ReplaceAll(strings.TrimSpace(*sa.TTMLOrigin), " ", ",")
 		sa.WebVTTScroll = "up"
-		//cue settings
+		// cue settings
 		coordinates := strings.Split(*sa.TTMLOrigin, " ")
 		if len(coordinates) > 1 {
 			sa.WebVTTLine = coordinates[0]
@@ -487,6 +487,51 @@ func (s *Subtitles) ForceDuration(d time.Duration, addDummyItem bool) {
 	}
 }
 
+func (s *Subtitles) Segment(segmentationType string, segmentDuration float64, segmentDurations []float64) []*Subtitles {
+	if len(s.Items) == 0 {
+		return nil
+	}
+	s.Order()
+	var noOfSegs int
+	if segmentationType == "SPECIFIED" {
+		noOfSegs = len(segmentDurations)
+	} else {
+		totalDurationSecs := s.Items[len(s.Items)-1].EndAt.Seconds()
+		noOfSegs = int(math.Ceil(totalDurationSecs / segmentDuration))
+	}
+	if noOfSegs == 0 {
+		return nil
+	}
+	var subs []*Subtitles
+	itemIdx := 0
+	end := time.Duration(0)
+	for i := 0; i < noOfSegs; i++ {
+		sub := NewSubtitles()
+		sub.Regions = s.Regions
+		sub.Styles = s.Styles
+		sub.Metadata = s.Metadata
+		if segmentationType == "SPECIFIED" {
+			end = end + time.Duration(segmentDurations[i]*float64(time.Second))
+		} else {
+			end = time.Duration(float64(i+1) * segmentDuration * float64(time.Second))
+		}
+		if end <= s.Items[itemIdx].StartAt {
+		} else {
+			for itemIdx < len(s.Items) && end > s.Items[itemIdx].StartAt {
+				sub.Items = append(sub.Items, s.Items[itemIdx])
+				if s.Items[itemIdx].EndAt <= end {
+					itemIdx++
+				} else {
+					// handle same cue in multiple segment
+					break
+				}
+			}
+		}
+		subs = append(subs, sub)
+	}
+	return subs
+}
+
 // Fragment fragments subtitles with a specific fragment duration
 func (s *Subtitles) Fragment(f time.Duration) {
 	// Nothing to fragment

webvtt.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"os"
 	"regexp"
 	"sort"
 	"strconv"
@@ -385,18 +386,35 @@ func formatDurationWebVTT(i time.Duration) string {
 	return formatDuration(i, ".", 3)
 }
 
+// WriteToWebVTTFile writes subtitles in .vtt format
+func (s Subtitles) WriteToWebVTTFile(dst string, offset float64) error {
+	// Do not write anything if no subtitles
+	f, err := os.Create(dst)
+	if err != nil {
+		return err
+	}
+	return s.WriteToWebVTTWithSync(f, offset)
+}
+
 // WriteToWebVTT writes subtitles in .vtt format
 func (s Subtitles) WriteToWebVTT(o io.Writer) (err error) {
 	// Do not write anything if no subtitles
 	if len(s.Items) == 0 {
 		err = ErrNoSubtitlesToWrite
 		return
 	}
+	return s.WriteToWebVTTWithSync(o, 0)
+}
 
+// WriteToWebVTTWithSync writes subtitles in .vtt format
+func (s Subtitles) WriteToWebVTTWithSync(o io.Writer, offset float64) (err error) {
 	// Add header
 	var c []byte
-	c = append(c, []byte("WEBVTT\n\n")...)
-
+	if offset == 0 {
+		c = append(c, []byte("WEBVTT\n\n")...)
+	} else {
+		c = append(c, []byte(fmt.Sprintf("WEBVTT\n%s=MPEGTS:%d,LOCAL:00:00:00.000\n\n", webvttTimestampMap, int(offset*90000)))...)
+	}
 	var style []string
 	for _, s := range s.Styles {
 		if s.InlineStyle != nil {