Michael McGuinness
fee71e7740
Feature/postprocessing * tests * passtest * fixshorttests * mosttests * improvingbasedockerfile * testspeeds * testing * host * canparallel * clean * passfullsuite * singlepagemax * test * findfeatures * findstables * tbls * tablestoo * tablestoo * lateraltests * tableloc * cleanup * inlinetable * childids * cleanup * tests
110 lines
2.5 KiB
Go
110 lines
2.5 KiB
Go
package documenttext
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"log/slog"
|
|
"strings"
|
|
|
|
documenttypes "queryorchestration/internal/document/types"
|
|
|
|
"github.com/aws/aws-sdk-go-v2/service/textract"
|
|
"github.com/aws/aws-sdk-go-v2/service/textract/types"
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
type cell struct {
|
|
value string
|
|
row int
|
|
column int
|
|
}
|
|
|
|
func (s *Service) addTable(ctx context.Context, block types.Block, blockMap map[uuid.UUID]types.Block, builder *strings.Builder) error {
|
|
table := []cell{}
|
|
maxRow := 0
|
|
maxColumn := 0
|
|
|
|
for _, uid := range s.getChildIds(block) {
|
|
childBlock := blockMap[uid]
|
|
if childBlock.BlockType != types.BlockTypeCell {
|
|
continue
|
|
}
|
|
|
|
var str strings.Builder
|
|
for _, uid := range s.getChildIds(childBlock) {
|
|
if blockMap[uid].BlockType == types.BlockTypeWord {
|
|
if str.String() != "" {
|
|
str.WriteString(" ")
|
|
}
|
|
str.WriteString(*blockMap[uid].Text)
|
|
}
|
|
if maxRow < int(*childBlock.RowIndex) {
|
|
maxRow = int(*childBlock.RowIndex)
|
|
}
|
|
if maxColumn < int(*childBlock.ColumnIndex) {
|
|
maxColumn = int(*childBlock.ColumnIndex)
|
|
}
|
|
table = append(table, cell{
|
|
column: int(*childBlock.ColumnIndex) - 1,
|
|
row: int(*childBlock.RowIndex) - 1,
|
|
value: str.String(),
|
|
})
|
|
}
|
|
}
|
|
|
|
tableArr := make([][]string, maxRow)
|
|
for i := range tableArr {
|
|
tableArr[i] = make([]string, maxColumn)
|
|
}
|
|
|
|
for _, cell := range table {
|
|
tableArr[cell.row][cell.column] = cell.value
|
|
}
|
|
|
|
jsonData, err := json.Marshal(tableArr)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
builder.WriteString("-------Table Start--------\n")
|
|
builder.WriteString(string(jsonData))
|
|
builder.WriteString("\n-------Table End--------\n")
|
|
|
|
return nil
|
|
}
|
|
|
|
func (s *Service) getTables(ctx context.Context, document documenttypes.File, index int) ([]uuid.UUID, map[uuid.UUID]types.Block, error) {
|
|
pageBytes, err := document.GetPageAsPNG(ctx, index)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
res, err := s.cfg.GetTextractClient().AnalyzeDocument(ctx, &textract.AnalyzeDocumentInput{
|
|
Document: &types.Document{
|
|
Bytes: pageBytes,
|
|
},
|
|
FeatureTypes: []types.FeatureType{
|
|
types.FeatureTypeTables,
|
|
},
|
|
})
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
tables := []uuid.UUID{}
|
|
for _, block := range res.Blocks {
|
|
if block.BlockType == types.BlockTypeTable {
|
|
id, err := uuid.Parse(*block.Id)
|
|
if err != nil {
|
|
slog.Warn("error parsing block id", "error", err, "id", *block.Id)
|
|
continue
|
|
}
|
|
|
|
tables = append(tables, id)
|
|
}
|
|
}
|
|
blocks := s.getBlockMap(res.Blocks)
|
|
|
|
return tables, blocks, nil
|
|
}
|