jpegai

type ColorSpace uint8

type DecodeInfo struct {
	Width         int
	Height        int
	ColorSpace    ColorSpace
	BitDepth      int
	ModelID       string
	Quality       int
	LatentShape   *LatentShape
	HasHyperprior bool
}

type Decoder struct {
	// contains filtered or unexported fields
}

type DecoderNetwork interface {
	// Decode transforms latent representation back to image pixels.
	// Input shape: [batch, latent_channels, latent_height, latent_width]
	// Output shape: [batch, channels, height, width]
	Decode(latents []float32, shape *LatentShape) ([]float32, error)

	// GetModelInfo returns metadata about the decoder model.
	GetModelInfo() (*ModelMetadata, error)
}

type DecodingContext struct {
	// Header is the parsed header.
	Header *JPEGAIHeader

	// Latents is the decoded latent representation.
	Latents *LatentRepresentation

	// Hyperprior is the decoded hyperprior data (if present).
	Hyperprior *HyperpriorData

	// Position tracks the current read position.
	Position int

	// Iterations tracks decoding iterations for security.
	Iterations int
}

type EncoderNetwork interface {
	// Encode transforms input image data to latent representation.
	// Input shape: [batch, channels, height, width]
	// Output shape: [batch, latent_channels, latent_height, latent_width]
	Encode(input []float32, width, height, channels int) ([]float32, *LatentShape, error)

	// GetModelInfo returns metadata about the encoder model.
	GetModelInfo() (*ModelMetadata, error)
}

type EntropyMode uint8

type HyperpriorData struct {
	// Shape describes the hyperprior latent shape.
	Shape *LatentShape

	// Data is the encoded hyperprior data.
	Data []byte

	// ScaleParams contains scale parameters for latent distributions.
	ScaleParams []float32

	// MeanParams contains mean parameters for latent distributions.
	MeanParams []float32
}

type InferenceEngine interface {
	// Infer runs the decoder neural network on latent data.
	// Input is the quantized latent representation.
	// Output is the reconstructed pixel data.
	Infer(latentData []float32) ([]float32, error)

	// GetModelInfo returns metadata about the loaded model.
	GetModelInfo() (*ModelMetadata, error)
}

type JPEGAIHeader struct {
	// Version is the bitstream version number.
	Version uint16

	// Profile is the coding profile.
	Profile Profile

	// Width is the image width in pixels.
	Width int

	// Height is the image height in pixels.
	Height int

	// ColorSpace indicates the color space.
	ColorSpace ColorSpace

	// BitDepth is the bit depth per sample.
	BitDepth int

	// ModelID identifies the neural network model.
	ModelID string

	// Quality is the quality parameter (0-100).
	Quality int

	// LatentShape describes the latent tensor dimensions.
	LatentShape *LatentShape

	// HasHyperprior indicates if hyperprior data is present.
	HasHyperprior bool

	// HasSideInfo indicates if additional side information is present.
	HasSideInfo bool

	// Checksum is the header checksum for validation.
	Checksum uint32
}

type LatentCodestream struct {
	// Data is the raw entropy-coded bitstream.
	Data []byte

	// Shape describes the expected decoded shape.
	Shape *LatentShape

	// EntropyMode indicates the entropy coding method.
	EntropyMode EntropyMode

	// TableOffset is the offset to entropy coding tables.
	TableOffset int

	// DataOffset is the offset to actual coded data.
	DataOffset int
}

type LatentRepresentation struct {
	// Shape describes the tensor dimensions.
	Shape *LatentShape

	// Data contains the quantized latent values.
	// Stored as int16 after quantization.
	Data []int16

	// Scale is the quantization scale factor for dequantization.
	Scale float32

	// Offset is the quantization offset for dequantization.
	Offset float32
}

type LatentShape struct {
	// Batch is the batch dimension (typically 1).
	Batch int

	// Channels is the number of latent channels.
	Channels int

	// Height is the spatial height of the latent tensor.
	Height int

	// Width is the spatial width of the latent tensor.
	Width int
}

type ModelLoader interface {
	// LoadModel loads a neural network model from the specified path.
	// Returns metadata about the loaded model.
	LoadModel(path string) (*ModelMetadata, error)

	// UnloadModel releases the currently loaded model.
	UnloadModel() error

	// IsModelLoaded returns true if a model is currently loaded.
	IsModelLoaded() bool
}

type ModelMetadata struct {
	// ModelID is the unique identifier for the model architecture.
	ModelID string

	// ModelVersion is the version of the model.
	ModelVersion uint16

	// Architecture describes the network architecture type.
	Architecture string

	// InputChannels is the number of input channels expected.
	InputChannels int

	// LatentChannels is the number of channels in latent space.
	LatentChannels int

	// DownsampleFactor is the spatial downsampling factor.
	DownsampleFactor int

	// HyperLatentChannels is the number of hyperprior latent channels.
	HyperLatentChannels int

	// Precision indicates the numerical precision (e.g., "float32", "float16").
	Precision string

	// ModelSize is the size of the model in bytes (for validation).
	ModelSize int64
}

type ParseError struct {
	Offset  int    // Byte offset where the error occurred
	Message string // Human-readable error description
	Cause   error  // Underlying error if any
}

type Parser struct {
	// contains filtered or unexported fields
}

type Profile uint8