138 lines
3.7 KiB
Go
138 lines
3.7 KiB
Go
package roboflow
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
|
|
"gitea.miguelmuniz.com/rogueking/junk2jive-server/internal/logger"
|
|
"gitea.miguelmuniz.com/rogueking/junk2jive-server/internal/response"
|
|
"github.com/joho/godotenv"
|
|
)
|
|
|
|
// ImageRequest represents the incoming JSON request with image data
|
|
type ImageRequest struct {
|
|
Image string `json:"image"` // Base64 encoded image
|
|
}
|
|
|
|
// RoboflowResponse represents the structured response from Roboflow API
|
|
type RoboflowResponse struct {
|
|
Predictions []struct {
|
|
Class string `json:"class"`
|
|
Confidence float64 `json:"confidence"`
|
|
} `json:"predictions"`
|
|
}
|
|
|
|
// Service represents a Roboflow service
|
|
type Service struct {
|
|
apiKey string
|
|
}
|
|
|
|
|
|
func init() {
|
|
// Load .env file if it exists
|
|
err := godotenv.Load()
|
|
if err != nil {
|
|
logger.Debug("No .env file found or error loading it: ", err)
|
|
} else {
|
|
logger.Debug(".env file loaded successfully")
|
|
}
|
|
}
|
|
// NewService creates a new Roboflow service instance
|
|
func NewService() *Service {
|
|
apiKey := os.Getenv("ROBOFLOW_API_KEY")
|
|
if apiKey == "" {
|
|
logger.Error("ROBOFLOW_API_KEY environment variable not set")
|
|
}
|
|
return &Service{
|
|
apiKey: apiKey,
|
|
}
|
|
}
|
|
|
|
// HandleImageRequest processes the JSON image request
|
|
func HandleImageRequest(w http.ResponseWriter, r *http.Request) {
|
|
// Decode the JSON request
|
|
var req ImageRequest
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
response.RespondWithError(w, r, http.StatusBadRequest, "Invalid JSON request", err)
|
|
return
|
|
}
|
|
|
|
// Check if image was provided
|
|
if req.Image == "" {
|
|
response.RespondWithError(w, r, http.StatusBadRequest, "No image provided in request", nil)
|
|
return
|
|
}
|
|
|
|
// Decode base64 image
|
|
imageData, err := base64.StdEncoding.DecodeString(req.Image)
|
|
if err != nil {
|
|
response.RespondWithError(w, r, http.StatusBadRequest, "Invalid base64 image", err)
|
|
return
|
|
}
|
|
|
|
service := NewService()
|
|
|
|
// Process the image with Roboflow
|
|
responser, err := service.AnalyzeImage(imageData)
|
|
if err != nil {
|
|
response.RespondWithError(w, r, http.StatusInternalServerError, "Error processing image", err)
|
|
return
|
|
}
|
|
|
|
// Return JSON response
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(responser)
|
|
}
|
|
|
|
// AnalyzeImage sends image bytes to Roboflow API
|
|
func (s *Service) AnalyzeImage(imageData []byte) (*RoboflowResponse, error) {
|
|
// Base64 encode the image data
|
|
base64Data := base64.StdEncoding.EncodeToString(imageData)
|
|
|
|
// Create the request
|
|
url := fmt.Sprintf("https://serverless.roboflow.com/taco-puuof/1?api_key=%s", s.apiKey)
|
|
req, err := http.NewRequest(http.MethodPost, url, bytes.NewBufferString(base64Data))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create request: %w", err)
|
|
}
|
|
|
|
// Set appropriate headers
|
|
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
|
|
|
// Send the request
|
|
client := &http.Client{}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to send request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
// Check response status
|
|
if resp.StatusCode != http.StatusOK {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
// Parse the response
|
|
var response RoboflowResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
|
|
return nil, fmt.Errorf("failed to parse response: %w", err)
|
|
}
|
|
|
|
return &response, nil
|
|
}
|
|
|
|
// GetDetectedObjects extracts the class names from the response
|
|
func GetDetectedObjects(response *RoboflowResponse) []string {
|
|
var objects []string
|
|
for _, prediction := range response.Predictions {
|
|
objects = append(objects, prediction.Class)
|
|
}
|
|
return objects
|
|
}
|