path: root/main.go

package main

import (
	"crypto/md5"
	"flag"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

const (
	subExt          = ".vtt"
	ffCmdOut        = "data/ff_commands"
	timeSep         = "-->"
	metadataPath    = "data/metadata.json"
	metadataPathCSV = "data/metadata.tsv"
	segmentSize     = "00:08:00"
)

var (
	outdir = "/mnt/nvme/data"
)

type Utterance struct {
	LeftTime  string
	RightTime string
	Text      string
	OutPath   string
	FD        *FileData
}

type FileData struct {
	VttPath   string
	AudioPath string
	AudioBase string
}

func NewFileData(vttPath string) *FileData {
	fd := &FileData{
		VttPath:   vttPath,
		AudioPath: strings.Trim(vttPath, subExt),
	}
	fd.AudioBase = path.Base(fd.AudioPath)
	return fd
}

func keysToSlice(req map[string]struct{}) []string {
	resp := []string{}
	for k := range req {
		resp = append(resp, k)
	}
	return resp
}

func mapToCSV(req map[string]string) [][]string {
	resp := [][]string{}
	for k, v := range req {
		resp = append(resp, []string{k, v})
	}
	return resp
}

func hashStr(req string) string {
	return fmt.Sprintf("%x", md5.Sum([]byte(req)))
}

func fpathToOutname(req string) string {
	basename := filepath.Base(req)
	h := hashStr(basename)
	return path.Join(outdir, h+".opus")
}

func linesToUtterances(lines []string, fd *FileData) []*Utterance {
	resp := []*Utterance{}
	for i, line := range lines {
		if !strings.Contains(line, timeSep) {
			continue
		}
		// get times
		splitted := strings.Split(line, timeSep)

		u := &Utterance{
			Text:      lines[i+1],
			LeftTime:  strings.TrimSpace(splitted[0]),
			RightTime: strings.TrimSpace(splitted[1]),
			FD:        fd,
		}
		u.OutPath = fmt.Sprintf("%s/%s_%s_%s.opus", outdir, fd.AudioBase,
			u.LeftTime, u.RightTime)

		if !adequateTimes(u.LeftTime, u.RightTime) {
			continue
		}
		resp = append(resp, u)

	}
	return resp
}

func adequateTimes(left, right string) bool {
	if len(left) == len("06:18.240") {
		left = "00:" + left
	}
	if len(right) == len("06:18.240") {
		right = "00:" + right
	}
	lts, err := time.Parse("15:04:05.000", left)
	if err != nil {
		fmt.Println(left, right, err)
	}
	rts, err := time.Parse("15:04:05.000", right)
	if err != nil {
		fmt.Println(left, right, err)
	}
	if rts.After(lts) {
		return true
	}
	return false
}

func utterancesToFileTextMap(utterances []*Utterance) map[string]string {
	resp := make(map[string]string)
	for _, ut := range utterances {
		resp[path.Base(ut.OutPath)] = ut.Text
	}
	return resp
}

func oneFileRun(filepath string) []*Utterance {
	fd := NewFileData(filepath)

	lines := readLines(fd.VttPath)
	utterances := linesToUtterances(lines, fd)

	return utterances
}

func dirRun(dirpath string) []*Utterance {
	resp := []*Utterance{}
	vttFiles := getFileList(dirpath, subExt)
	for _, vtt := range vttFiles {
		utterances := oneFileRun(vtt)
		resp = append(resp, utterances...)
	}
	return resp
}

func getFileList(dirpath string, filter string) []string {
	resp := []string{}
	err := filepath.Walk(dirpath,
		func(path string, info os.FileInfo, err error) error {
			if err != nil {
				fmt.Printf("prevent panic by handling failure accessing a path %q: %v\n", path, err)
				return err
			}
			if strings.HasSuffix(info.Name(), filter) {
				resp = append(resp, path)
			}
			return nil
		})
	if err != nil {
		fmt.Println(err)
	}
	return resp
}

func equalSliceRun(dirpath string) {
	auFiles := getFileList(dirpath, "opus")

	fQueue := make(chan string, len(auFiles))
	for _, fpath := range auFiles {
		fQueue <- fpath
	}

	workers := 20
	for i := 0; i < workers; i++ {
		go cutterQueue(fQueue, i)
	}

	for {
		if len(fQueue) == 0 {
			fmt.Println("empty queue: work is done")
			break
		}
		time.Sleep(1 * time.Second)
	}
	close(fQueue)
	fmt.Println("queue closed")
}

func main() {
	vttFilepath := flag.String("f", "", "path to a vtt file")
	vttDir := flag.String("d", "", "path to a vtt dir")
	outDirArg := flag.String("o", "", "output dir")
	sliceAudioDir := flag.String("slice-audio-dir", "",
		"for equal segmentation only without subs")
	flag.Parse()

	if outDirArg != nil && *outDirArg != "" {
		outdir = *outDirArg
	}

	utterances := []*Utterance{}
	if vttDir != nil && *vttDir != "" {
		utterances = dirRun(*vttDir)
	} else if vttFilepath != nil && *vttFilepath != "" {
		utterances = oneFileRun(*vttFilepath)
	} else if sliceAudioDir != nil && *sliceAudioDir != "" {
		equalSliceRun(*sliceAudioDir)
		return
	} else {
		fmt.Println("no flags provided;")
		return
	}

	fmt.Println("sum of utterances:", len(utterances))

	if err := os.MkdirAll(outdir, 0755); err != nil {
		panic(err)
	}

	filteredUtterances := []*Utterance{}

	// === concurrent ===
	queue := make(chan *Utterance, len(utterances))
	geshaft := make(chan *Utterance, len(utterances))
	done := make(chan bool, 1)

	for _, ut := range utterances {
		if _, err := os.Stat(ut.OutPath); os.IsNotExist(err) {
			queue <- ut
			// if err := cutoutClipAndTranscode(ut); err == nil {
			// 	filteredUtterances = append(filteredUtterances, ut)
			// }
		} else if err == nil {
			geshaft <- ut
		}
	}

	fmt.Println("filled queue, starting workers")

	workers := 100
	for i := 0; i < workers; i++ {
		go worker(queue, done, i, geshaft)
	}

	for {
		if len(queue) == 0 {
			fmt.Println("empty queue: sending true to done")
			done <- true
			// give time for workers to finish
			fmt.Println("sleeping for few seconds")
			time.Sleep(5 * time.Second)
			fmt.Println("closing queue")
			close(queue)
			fmt.Println("breaking from the loop")
			break
		}
		time.Sleep(1 * time.Second)
	}

	close(geshaft)
	for ut := range geshaft {
		filteredUtterances = append(filteredUtterances, ut)
	}
	newMeta := utterancesToFileTextMap(filteredUtterances)
	writeCSV(mapToCSV(newMeta))
}