github.com/instill-ai/component@v0.16.0-beta/pkg/operator/image/v0/main.go (about)

     1  //go:generate compogen readme --operator ./config ./README.mdx
     2  package image
     3  
     4  import (
     5  	"bytes"
     6  	"encoding/base64"
     7  	"fmt"
     8  	"image"
     9  	"sync"
    10  
    11  	_ "embed" // embed
    12  	_ "image/gif"
    13  	_ "image/jpeg"
    14  	_ "image/png"
    15  
    16  	"go.uber.org/zap"
    17  	"google.golang.org/protobuf/types/known/structpb"
    18  
    19  	"github.com/instill-ai/component/pkg/base"
    20  )
    21  
    22  var (
    23  	//go:embed config/definition.json
    24  	definitionJSON []byte
    25  	//go:embed config/tasks.json
    26  	tasksJSON []byte
    27  	once      sync.Once
    28  	op        *operator
    29  )
    30  
    31  // Operator is the derived operator
    32  type operator struct {
    33  	base.BaseOperator
    34  }
    35  
    36  // Execution is the derived execution
    37  type execution struct {
    38  	base.BaseOperatorExecution
    39  }
    40  
    41  // Init initializes the operator
    42  func Init(l *zap.Logger, u base.UsageHandler) *operator {
    43  	once.Do(func() {
    44  		op = &operator{
    45  			BaseOperator: base.BaseOperator{
    46  				Logger:       l,
    47  				UsageHandler: u,
    48  			},
    49  		}
    50  		err := op.LoadOperatorDefinition(definitionJSON, tasksJSON, nil)
    51  		if err != nil {
    52  			panic(err)
    53  		}
    54  	})
    55  	return op
    56  }
    57  
    58  func (o *operator) CreateExecution(sysVars map[string]any, task string) (*base.ExecutionWrapper, error) {
    59  	return &base.ExecutionWrapper{Execution: &execution{
    60  		BaseOperatorExecution: base.BaseOperatorExecution{Operator: o, SystemVariables: sysVars, Task: task},
    61  	}}, nil
    62  }
    63  
    64  // Execute executes the derived execution
    65  func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
    66  	outputs := []*structpb.Struct{}
    67  	var base64ByteImg []byte
    68  	for _, input := range inputs {
    69  
    70  		b, err := base64.StdEncoding.DecodeString(base.TrimBase64Mime(input.Fields["image"].GetStringValue()))
    71  		if err != nil {
    72  			return nil, err
    73  		}
    74  
    75  		img, _, err := image.Decode(bytes.NewReader(b))
    76  		if err != nil {
    77  			return nil, err
    78  		}
    79  
    80  		switch e.Task {
    81  		case "TASK_DRAW_CLASSIFICATION":
    82  			base64ByteImg, err = drawClassification(img, input.Fields["category"].GetStringValue(), input.Fields["score"].GetNumberValue())
    83  			if err != nil {
    84  				return nil, err
    85  			}
    86  		case "TASK_DRAW_DETECTION":
    87  			base64ByteImg, err = drawDetection(img, input.Fields["objects"].GetListValue().Values)
    88  			if err != nil {
    89  				return nil, err
    90  			}
    91  		case "TASK_DRAW_KEYPOINT":
    92  			base64ByteImg, err = drawKeypoint(img, input.Fields["objects"].GetListValue().Values)
    93  			if err != nil {
    94  				return nil, err
    95  			}
    96  		case "TASK_DRAW_OCR":
    97  			base64ByteImg, err = drawOCR(img, input.Fields["objects"].GetListValue().Values)
    98  			if err != nil {
    99  				return nil, err
   100  			}
   101  		case "TASK_DRAW_INSTANCE_SEGMENTATION":
   102  			base64ByteImg, err = drawInstanceSegmentation(img, input.Fields["objects"].GetListValue().Values)
   103  			if err != nil {
   104  				return nil, err
   105  			}
   106  		case "TASK_DRAW_SEMANTIC_SEGMENTATION":
   107  			base64ByteImg, err = drawSemanticSegmentation(img, input.Fields["stuffs"].GetListValue().Values)
   108  			if err != nil {
   109  				return nil, err
   110  			}
   111  		default:
   112  			return nil, fmt.Errorf("not supported task: %s", e.Task)
   113  		}
   114  
   115  		output := structpb.Struct{Fields: make(map[string]*structpb.Value)}
   116  
   117  		output.Fields["image"] = &structpb.Value{
   118  			Kind: &structpb.Value_StringValue{
   119  				StringValue: fmt.Sprintf("data:image/jpeg;base64,%s", string(base64ByteImg)),
   120  			},
   121  		}
   122  
   123  		outputs = append(outputs, &output)
   124  	}
   125  	return outputs, nil
   126  }