query-orchestration/cmd/textractGenerator_test/textdetection.go

//go:build aws

//go:generate go run -tags=aws ./textdetection.go
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	documenttext "queryorchestration/internal/document/text"
	documenttypes "queryorchestration/internal/document/types"
	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/serviceconfig/aws"
	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/serviceconfig/queue/querysync"
	"queryorchestration/internal/serviceconfig/textract"
	"sync"
	"time"

	awstextract "github.com/aws/aws-sdk-go-v2/service/textract"
	"github.com/aws/aws-sdk-go-v2/service/textract/types"
)

type Config struct {
	serviceconfig.BaseConfig
	textract.TextractConfig
	querysync.QuerySyncConfig
	objectstore.ObjectStoreConfig
}

func main() {
	ctx := context.Background()

	if len(os.Args) != 3 {
		fmt.Println("Error: expected at least two arguments")
		fmt.Println("Usage:", os.Args[0], "<sample_name>", "<file_name>")
		os.Exit(1)
	}

	sampleName := os.Args[1]
	name := os.Args[2]

	slog.Info("arguments", "sample_name", sampleName, "file_name", name)

	cfg := &Config{}
	profile := aws.Profile(os.Getenv("AWS_PROFILE"))
	if profile == "" {
		slog.Error("profile not available")
		os.Exit(1)
	}

	err := cfg.SetTextractClientWithProfile(ctx, profile)
	if err != nil {
		slog.Error("error creating client", "error", err)
		os.Exit(1)
	}

	root, err := os.OpenRoot("assets")
	if err != nil {
		slog.Error("error opening root", "error", err)
		os.Exit(1)
	}

	root, err = root.OpenRoot(sampleName)
	if err != nil {
		slog.Error("error opening sample", "error", err)
		os.Exit(1)
	}

	pdfFile, err := root.Open(fmt.Sprintf("%s.pdf", name))
	if err != nil {
		slog.Error("error opening pdf file", "error", err)
		os.Exit(1)
	}

	pdf, err := documenttypes.NewPDFFromReader(pdfFile)
	if err != nil {
		slog.Error("error creating pdf file", "error", err)
		os.Exit(1)
	}

	textSvc := documenttext.New(cfg)

	count, err := pdf.GetPageCount(ctx)
	if err != nil {
		slog.Error("error getting page count", "error", err)
		os.Exit(1)
	}

	slog.Info("detecting text", "page_count", count)

	startTime := time.Now()
	allResults := map[string][]awstextract.AnalyzeDocumentOutput{}
	allResults["base"] = make([]awstextract.AnalyzeDocumentOutput, count)
	var wg sync.WaitGroup
	for i := range count {
		wg.Add(1)
		go func() {
			slog.Info("getting page", "index", i)

			result, err := textSvc.GetBasePage(ctx, pdf, i)
			if err != nil {
				slog.Error("error detecting text", "error", err)
				os.Exit(1)
			}

			allResults["base"][i] = result
			wg.Done()
		}()
	}
	wg.Wait()

	allResults["table"] = []awstextract.AnalyzeDocumentOutput{}
	for i, result := range allResults["base"] {
		hasTable := false
		for _, block := range result.Blocks {
			if block.BlockType == types.BlockTypeLayoutTable {
				hasTable = true
			}
		}

		if hasTable {
			slog.Info("getting table", "index", i)
			pageBytes, err := pdf.GetPageAsPNG(ctx, i)
			if err != nil {
				slog.Error("error getting page", "error", err)
				os.Exit(1)
			}

			res, err := cfg.GetTextractClient().AnalyzeDocument(ctx, &awstextract.AnalyzeDocumentInput{
				Document: &types.Document{
					Bytes: pageBytes,
				},
				FeatureTypes: []types.FeatureType{
					types.FeatureTypeTables,
				},
			})
			if err != nil {
				slog.Error("error detecting table", "error", err)
				os.Exit(1)
			}
			allResults["table"] = append(allResults["table"], *res)
		}
	}

	generatedBase := "generated"

	if _, err := root.Stat(generatedBase); os.IsNotExist(err) {
		err := root.Mkdir(generatedBase, 0755)
		if err != nil {
			slog.Error("error creating directory", "error", err)
			os.Exit(1)
		}

		slog.Info("directory created", "directory", generatedBase)
	}

	jsonContent, err := json.MarshalIndent(allResults, "", "\t")
	if err != nil {
		slog.Error("error marshaling json", "error", err)
		os.Exit(1)
	}

	outputFilename := fmt.Sprintf("%s/%s.gen", generatedBase, name)
	file, err := root.Create(outputFilename)
	if err != nil {
		slog.Error("error creating file", "error", err)
		os.Exit(1)
	}
	defer file.Close()

	_, err = file.Write(jsonContent)
	if err != nil {
		slog.Error("error writing file", "error", err)
		os.Exit(1)
	}

	slog.Info("extracted text", "output", outputFilename, "time", time.Since(startTime))
}