hatch-surf/internal/face/rekognition.go

package face

import (
	"context"
	"fmt"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"bytes"
	"io"
	"mime/multipart"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/rekognition"
	"github.com/aws/aws-sdk-go-v2/service/rekognition/types"
)

// RekognitionClient wraps AWS Rekognition operations
type RekognitionClient struct {
	client       *rekognition.Client
	collectionID string
}

// FaceDetection represents a detected face
type FaceDetection struct {
	Confidence float64     `json:"confidence"`
	BoundingBox BoundingBox `json:"boundingBox"`
}

// BoundingBox represents face location in image
type BoundingBox struct {
	Left   float64 `json:"left"`
	Top    float64 `json:"top"`
	Width  float64 `json:"width"`
	Height float64 `json:"height"`
}

// FaceMatch represents a matched face
type FaceMatch struct {
	FaceID          string  `json:"faceId"`
	ExternalImageID string  `json:"externalImageId"`
	Similarity      float64 `json:"similarity"`
	BoundingBox     BoundingBox `json:"boundingBox"`
	PhotoID         string  `json:"photoId,omitempty"`
	PhotoURL        string  `json:"photoUrl,omitempty"`
	ThumbnailURL    string  `json:"thumbnailUrl,omitempty"`
}

// NewRekognitionClient creates a new Rekognition client
func NewRekognitionClient(cfg aws.Config, collectionID string) *RekognitionClient {
	return &RekognitionClient{
		client:       rekognition.NewFromConfig(cfg),
		collectionID: collectionID,
	}
}

// CreateCollection creates a new face collection
func (r *RekognitionClient) CreateCollection(ctx context.Context, collectionName string) error {
	_, err := r.client.CreateCollection(ctx, &rekognition.CreateCollectionInput{
		CollectionId: &collectionName,
	})
	if err != nil {
		return fmt.Errorf("failed to create collection: %w", err)
	}
	return nil
}

// DeleteCollection deletes a face collection
func (r *RekognitionClient) DeleteCollection(ctx context.Context, collectionName string) error {
	_, err := r.client.DeleteCollection(ctx, &rekognition.DeleteCollectionInput{
		CollectionId: &collectionName,
	})
	if err != nil {
		return fmt.Errorf("failed to delete collection: %w", err)
	}
	return nil
}

// DetectFaces detects faces in an image
func (r *RekognitionClient) DetectFaces(ctx context.Context, imageBytes []byte) ([]FaceDetection, error) {
	input := &rekognition.DetectFacesInput{
		Image: &types.Image{
			Bytes: imageBytes,
		},
		Attributes: []types.Attribute{
			types.AttributeAll,
		},
	}

	result, err := r.client.DetectFaces(ctx, input)
	if err != nil {
		return nil, fmt.Errorf("failed to detect faces: %w", err)
	}

	var detections []FaceDetection
	for _, face := range result.FaceDetails {
		if face.BoundingBox == nil {
			continue
		}

		detection := FaceDetection{
			Confidence: float64(*face.Confidence),
			BoundingBox: BoundingBox{
				Left:   float64(*face.BoundingBox.Left),
				Top:    float64(*face.BoundingBox.Top),
				Width:  float64(*face.BoundingBox.Width),
				Height: float64(*face.BoundingBox.Height),
			},
		}
		detections = append(detections, detection)
	}

	return detections, nil
}

// IndexFace indexes a face in the collection
func (r *RekognitionClient) IndexFace(ctx context.Context, imageBytes []byte, photoID string) (*FaceDetection, error) {
	input := &rekognition.IndexFacesInput{
		CollectionId: &r.collectionID,
		Image: &types.Image{
			Bytes: imageBytes,
		},
		ExternalImageId: &photoID,
		MaxFaces:        aws.Int32(1),
		QualityFilter:   types.QualityFilterAuto,
	}

	result, err := r.client.IndexFaces(ctx, input)
	if err != nil {
		return nil, fmt.Errorf("failed to index face: %w", err)
	}

	if len(result.FaceRecords) == 0 {
		return nil, fmt.Errorf("no faces found in image")
	}

	faceRecord := result.FaceRecords[0]
	if faceRecord.FaceDetail == nil || faceRecord.FaceDetail.BoundingBox == nil {
		return nil, fmt.Errorf("no face detail available")
	}

	return &FaceDetection{
		Confidence: float64(*faceRecord.Face.Confidence),
		BoundingBox: BoundingBox{
			Left:   float64(*faceRecord.FaceDetail.BoundingBox.Left),
			Top:    float64(*faceRecord.FaceDetail.BoundingBox.Top),
			Width:  float64(*faceRecord.FaceDetail.BoundingBox.Width),
			Height: float64(*faceRecord.FaceDetail.BoundingBox.Height),
		},
	}, nil
}

// SearchFacesByImage searches for similar faces
func (r *RekognitionClient) SearchFacesByImage(ctx context.Context, imageBytes []byte, maxFaces int, threshold float64) ([]FaceMatch, error) {
	input := &rekognition.SearchFacesByImageInput{
		CollectionId: &r.collectionID,
		Image: &types.Image{
			Bytes: imageBytes,
		},
		MaxFaces:          aws.Int32(int32(maxFaces)),
		FaceMatchThreshold: aws.Float32(float32(threshold)),
	}

	result, err := r.client.SearchFacesByImage(ctx, input)
	if err != nil {
		return nil, fmt.Errorf("failed to search faces: %w", err)
	}

	var matches []FaceMatch
	for _, match := range result.FaceMatches {
		if match.Face == nil || match.Face.BoundingBox == nil {
			continue
		}

		faceMatch := FaceMatch{
			FaceID:          aws.ToString(match.Face.FaceId),
			ExternalImageID: aws.ToString(match.Face.ExternalImageId),
			Similarity:      float64(aws.ToFloat32(match.Similarity)),
			BoundingBox: BoundingBox{
				Left:   float64(*match.Face.BoundingBox.Left),
				Top:    float64(*match.Face.BoundingBox.Top),
				Width:  float64(*match.Face.BoundingBox.Width),
				Height: float64(*match.Face.BoundingBox.Height),
			},
		}
		matches = append(matches, faceMatch)
	}

	return matches, nil
}

// ExtractImageBytes extracts image bytes from a multipart file
func ExtractImageBytes(file multipart.File) ([]byte, error) {
	// Read all bytes
	data, err := io.ReadAll(file)
	if err != nil {
		return nil, fmt.Errorf("failed to read file: %w", err)
	}

	// Validate it's an image
	_, _, err = image.DecodeConfig(bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("invalid image format: %w", err)
	}

	return data, nil
}

// NewRekognitionClientFromEnv creates a client from environment variables
func NewRekognitionClientFromEnv(collectionID string) (*RekognitionClient, error) {
	cfg, err := config.LoadDefaultConfig(context.Background())
	if err != nil {
		return nil, fmt.Errorf("failed to load AWS config: %w", err)
	}

	return NewRekognitionClient(cfg, collectionID), nil
}