[AI] Add support for Object Detection pipeline

common/util.go

@@ -462,6 +462,21 @@ func ParseEthAddr(strJsonKey string) (string, error) {
 	return "", errors.New("Error parsing address from keyfile")
 }
 
+func GetInputVideoInfo(video types.File) (ffmpeg.MediaFormatInfo, error) {
+	bytearr, _ := video.Bytes()
+	_, mediaFormat, err := ffmpeg.GetCodecInfoBytes(bytearr)
+	if err != nil {
+		return ffmpeg.MediaFormatInfo{}, errors.New("Error getting codec info")
+	}
+
+	duration := int64(mediaFormat.DurSecs)
+	if duration <= 0 {
+		return ffmpeg.MediaFormatInfo{}, errors.New("video duration calculation failed")
+	}
+
+	return mediaFormat, nil
+}
+
 // CalculateAudioDuration calculates audio file duration using the lpms/ffmpeg package.
 func CalculateAudioDuration(audio types.File) (int64, error) {
 	read, err := audio.Reader()

core/ai.go

@@ -28,6 +28,7 @@ type AI interface {
 	ImageToText(context.Context, worker.GenImageToTextMultipartRequestBody) (*worker.ImageToTextResponse, error)
 	TextToSpeech(context.Context, worker.GenTextToSpeechJSONRequestBody) (*worker.AudioResponse, error)
 	LiveVideoToVideo(context.Context, worker.GenLiveVideoToVideoJSONRequestBody) (*worker.LiveVideoToVideoResponse, error)
+	ObjectDetection(context.Context, worker.GenObjectDetectionMultipartRequestBody) (*worker.ObjectDetectionResponse, error)
 	Warm(context.Context, string, string, worker.RunnerEndpoint, worker.OptimizationFlags) error
 	Stop(context.Context) error
 	HasCapacity(string, string) bool

core/ai_test.go

@@ -667,6 +667,15 @@ func (a *stubAIWorker) LiveVideoToVideo(ctx context.Context, req worker.GenLiveV
 	return &worker.LiveVideoToVideoResponse{}, nil
 }
 
+func (a *stubAIWorker) ObjectDetection(ctx context.Context, req worker.GenObjectDetectionMultipartRequestBody) (*worker.ObjectDetectionResponse, error) {
+	return &worker.ObjectDetectionResponse{
+		Video:            {Url: "http://example.com/frames1.mp4"},
+		ConfidenceScores: "confidence_scores",
+		Labels:           "labels",
+		DetectionBoxes:   "detection_boxes",
+		DetectionPts:     "detection_pts"}, nil
+}
+
 func (a *stubAIWorker) Warm(ctx context.Context, arg1, arg2 string, endpoint worker.RunnerEndpoint, flags worker.OptimizationFlags) error {
 	return nil
 }

core/ai_worker.go

@@ -465,6 +465,21 @@
 		}
 		resp.Audio.Url = osUrl
 
+		results = resp
+	case worker.ObjectDetectionResponse:
+		if resp.Video.Url != "" {
+			err := worker.ReadVideoB64DataUrl(resp.Video.Url, &buf)
+			if err != nil {
+				return nil, err
+			}
+
+			osUrl, err := storage.OS.SaveData(ctx, fileName, bytes.NewBuffer(buf.Bytes()), nil, 0)
+			if err != nil {
+				return nil, err
+			}
+			resp.Video.Url = osUrl
+		}
+
 		results = resp
 	}
 
@@ -510,6 +525,19 @@
 		delete(results.Files, fileName)
 
 		results.Results = resp
+	case worker.ObjectDetectionResponse:
+		if resp.Video.Url != "" {
+			fileName := resp.Video.Url
+			osUrl, err := storage.OS.SaveData(ctx, fileName, bytes.NewReader(results.Files[fileName]), nil, 0)
+			if err != nil {
+				return nil, err
+			}
+
+			resp.Video.Url = osUrl
+			delete(results.Files, fileName)
+
+			results.Results = resp
+		}
 	}
 
 	// no file response to save, response is text
@@ -884,6 +912,50 @@
 	return res.Results, nil
 }
 
+func (orch *orchestrator) ObjectDetection(ctx context.Context, requestID string, req worker.GenObjectDetectionMultipartRequestBody) (interface{}, error) {
+	// local AIWorker processes job if combined orchestrator/ai worker
+	if orch.node.AIWorker != nil {
+		workerResp, err := orch.node.ObjectDetection(ctx, req)
+		if err == nil {
+			return orch.node.saveLocalAIWorkerResults(ctx, *workerResp, requestID, "video/mp4")
+		} else {
+			clog.Errorf(ctx, "Error processing with local ai worker err=%q", err)
+			if monitor.Enabled {
+				monitor.AIResultSaveError(ctx, "object-detection", *req.ModelId, string(monitor.SegmentUploadErrorUnknown))
+			}
+			return nil, err
+		}
+	}
+
+	// remote ai worker proceses job
+	videoBytes, err := req.Video.Bytes()
+	if err != nil {
+		return nil, err
+	}
+
+	inputUrl, err := orch.SaveAIRequestInput(ctx, requestID, videoBytes)
+	if err != nil {
+		return nil, err
+	}
+	req.Video.InitFromBytes(nil, "")
+
+	res, err := orch.node.AIWorkerManager.Process(ctx, requestID, "object-detection", *req.ModelId, inputUrl, AIJobRequestData{Request: req, InputUrl: inputUrl})
+	if err != nil {
+		return nil, err
+	}
+
+	res, err = orch.node.saveRemoteAIWorkerResults(ctx, res, requestID)
+	if err != nil {
+		clog.Errorf(ctx, "Error saving remote ai result err=%q", err)
+		if monitor.Enabled {
+			monitor.AIResultSaveError(ctx, "object-detection", *req.ModelId, string(monitor.SegmentUploadErrorUnknown))
+		}
+		return nil, err
+	}
+
+	return res.Results, nil
+}
+
 // only used for sending work to remote AI worker
 func (orch *orchestrator) SaveAIRequestInput(ctx context.Context, requestID string, fileData []byte) (string, error) {
 	node := orch.node
@@ -1062,7 +1134,11 @@
 	return n.AIWorker.LiveVideoToVideo(ctx, req)
 }
 
-// transcodeFrames converts a series of image URLs into a video segment for the image-to-video pipeline.
+func (n *LivepeerNode) ObjectDetection(ctx context.Context, req worker.GenObjectDetectionMultipartRequestBody) (*worker.ObjectDetectionResponse, error) {
+	return n.AIWorker.ObjectDetection(ctx, req)
+}
+
+// transcodeFrames converts a series of image URLs into a video segment for the image-to-video and object-detection pipeline.
 func (n *LivepeerNode) transcodeFrames(ctx context.Context, sessionID string, urls []string, inProfile ffmpeg.VideoProfile, outProfile ffmpeg.VideoProfile) *TranscodeResult {
 	ctx = clog.AddOrchSessionID(ctx, sessionID)
 

core/capabilities.go

@@ -83,6 +83,7 @@ const (
 	Capability_ImageToText                Capability = 34
 	Capability_LiveVideoToVideo           Capability = 35
 	Capability_TextToSpeech               Capability = 36
+	Capability_ObjectDetection            Capability = 37
 )
 
 var CapabilityNameLookup = map[Capability]string{
@@ -124,6 +125,7 @@ var CapabilityNameLookup = map[Capability]string{
 	Capability_ImageToText:                "Image to text",
 	Capability_LiveVideoToVideo:           "Live video to video",
 	Capability_TextToSpeech:               "Text to speech",
+	Capability_ObjectDetection:            "Object detection",
 }
 
 var CapabilityTestLookup = map[Capability]CapabilityTest{
@@ -217,6 +219,7 @@ func OptionalCapabilities() []Capability {
 		Capability_SegmentAnything2,
 		Capability_ImageToText,
 		Capability_TextToSpeech,
+		Capability_ObjectDetection,
 	}
 }
 

go.mod

@@ -257,3 +257,5 @@ require (
 	lukechampine.com/blake3 v1.2.1 // indirect
 	rsc.io/tmplfunc v0.0.3 // indirect
 )
+
+replace github.com/livepeer/ai-worker => github.com/RUFFY-369/ai-worker v0.8.1-0.20241102154421-60e5d350c2df

go.sum

@@ -63,6 +63,8 @@ github.com/Masterminds/semver/v3 v3.2.1 h1:RN9w6+7QoMeJVGyfmbcgs28Br8cvmnucEXnY0
 github.com/Masterminds/semver/v3 v3.2.1/go.mod h1:qvl/7zhW3nngYb5+80sSMF+FG2BjYrf8m9wsX0PNOMQ=
 github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
 github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
+github.com/RUFFY-369/ai-worker v0.8.1-0.20241102154421-60e5d350c2df h1:UL+t2GVDDk20eypLFcbbqiO95OmlyzjARCkBUbvVNmc=
+github.com/RUFFY-369/ai-worker v0.8.1-0.20241102154421-60e5d350c2df/go.mod h1:GjQuPmz69UO53WVtqzB9Ygok5MmKCGNuobbfMXH7zgw=
 github.com/RaveNoX/go-jsoncommentstrip v1.0.0/go.mod h1:78ihd09MekBnJnxpICcwzCMzGrKSKYe4AqU6PDYYpjk=
 github.com/Shopify/goreferrer v0.0.0-20181106222321-ec9c9a553398/go.mod h1:a1uqRtAwp2Xwc6WNPJEufxJ7fx3npB4UV/JOLmbu5I0=
 github.com/StackExchange/wmi v1.2.1 h1:VIkavFPXSjcnS+O8yTq7NI32k0R5Aj+v39y29VYDOSA=

server/ai_http.go

@@ -71,6 +71,7 @@ func startAIServer(lp *lphttp) error {
 	lp.transRPC.Handle("/image-to-text", oapiReqValidator(aiHttpHandle(lp, multipartDecoder[worker.GenImageToTextMultipartRequestBody])))
 	lp.transRPC.Handle("/text-to-speech", oapiReqValidator(aiHttpHandle(lp, jsonDecoder[worker.GenTextToSpeechJSONRequestBody])))
 	lp.transRPC.Handle("/live-video-to-video", oapiReqValidator(lp.StartLiveVideoToVideo()))
+	lp.transRPC.Handle("/object-detection", oapiReqValidator(aiHttpHandle(&lp, multipartDecoder[worker.GenObjectDetectionMultipartRequestBody])))
 	// Additionally, there is the '/aiResults' endpoint registered in server/rpc.go
 
 	return nil
@@ -470,6 +471,20 @@ func handleAIRequest(ctx context.Context, w http.ResponseWriter, r *http.Request
 		// TTS pricing is typically in characters, including punctuation.
 		words := utf8.RuneCountInString(*v.Text)
 		outPixels = int64(1000 * words)
+	case worker.GenObjectDetectionMultipartRequestBody:
+		pipeline = "object-detection"
+		cap = core.Capability_ObjectDetection
+		modelID = *v.ModelId
+		mediaFormat, err := common.GetInputVideoInfo(v.Video)
+		if err != nil {
+			respondWithError(w, err.Error(), http.StatusBadRequest)
+		}
+
+		submitFn = func(ctx context.Context) (interface{}, error) {
+			return orch.ObjectDetection(ctx, requestID, v)
+		}
+		// Calculate the output pixels using the video profile
+		outPixels = int64(mediaFormat.Width) * int64(mediaFormat.Height) * int64(mediaFormat.FPS) * mediaFormat.DurSecs
 	default:
 		respondWithError(w, "Unknown request type", http.StatusBadRequest)
 		return
@@ -575,6 +590,8 @@ func handleAIRequest(ctx context.Context, w http.ResponseWriter, r *http.Request
 			latencyScore = CalculateImageToTextLatencyScore(took, outPixels)
 		case worker.GenTextToSpeechJSONRequestBody:
 			latencyScore = CalculateTextToSpeechLatencyScore(took, outPixels)
+		case worker.GenObjectDetectionMultipartRequestBody:
+			latencyScore = CalculateObjectDetectionLatencyScore(took, outPixels)
 		}
 
 		var pricePerAIUnit float64
@@ -786,6 +803,16 @@ func parseMultiPartResult(body io.Reader, boundary string, pipeline string) core
 				}
 			case "text-to-speech":
 				var parsedResp worker.AudioResponse
+				err := json.Unmarshal(body, &parsedResp)
+				if err != nil {
+					glog.Error("Error getting results json:", err)
+					wkrResult.Err = err
+					break
+				}
+				results = parsedResp
+			case "object-detection":
+				var parsedResp worker.ObjectDetectionResponse
+
 				err := json.Unmarshal(body, &parsedResp)
 				if err != nil {
 					glog.Error("Error getting results json:", err)

server/ai_mediaserver.go

@@ -81,6 +81,7 @@ func startAIMediaServer(ls *LivepeerServer) error {
 	ls.HTTPMux.Handle("/segment-anything-2", oapiReqValidator(aiMediaServerHandle(ls, multipartDecoder[worker.GenSegmentAnything2MultipartRequestBody], processSegmentAnything2)))
 	ls.HTTPMux.Handle("/image-to-text", oapiReqValidator(aiMediaServerHandle(ls, multipartDecoder[worker.GenImageToTextMultipartRequestBody], processImageToText)))
 	ls.HTTPMux.Handle("/text-to-speech", oapiReqValidator(aiMediaServerHandle(ls, jsonDecoder[worker.GenTextToSpeechJSONRequestBody], processTextToSpeech)))
+	ls.HTTPMux.Handle("/object-detection", oapiReqValidator(aiMediaServerHandle(ls, multipartDecoder[worker.GenObjectDetectionMultipartRequestBody], processObjectDetection)))
 
 	// This is called by the media server when the stream is ready
 	ls.HTTPMux.Handle("/live/video-to-video/{stream}/start", ls.StartLiveVideo())