mirror of
https://github.com/songquanpeng/one-api.git
synced 2025-09-18 17:46:37 +08:00
671fe78e44
- Add support for getting base64-encoded images via openAI's image_url. - Add `context` as a parameter for the function `LogError`. - Handle the error from `image.GetImageFromUrl` by logging it. - Convert the role to `user` if it is `system` and add a dummy model message to make Gemini happy.
105 lines
2.1 KiB
Go
105 lines
2.1 KiB
Go
package image
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/base64"
|
|
"image"
|
|
_ "image/gif"
|
|
_ "image/jpeg"
|
|
_ "image/png"
|
|
"net/http"
|
|
"regexp"
|
|
"strings"
|
|
"sync"
|
|
|
|
_ "golang.org/x/image/webp"
|
|
)
|
|
|
|
func IsImageUrl(url string) (bool, error) {
|
|
resp, err := http.Head(url)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if !strings.HasPrefix(resp.Header.Get("Content-Type"), "image/") {
|
|
return false, nil
|
|
}
|
|
return true, nil
|
|
}
|
|
|
|
func GetImageSizeFromUrl(url string) (width int, height int, err error) {
|
|
isImage, err := IsImageUrl(url)
|
|
if !isImage {
|
|
return
|
|
}
|
|
resp, err := http.Get(url)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
img, _, err := image.DecodeConfig(resp.Body)
|
|
if err != nil {
|
|
return
|
|
}
|
|
return img.Width, img.Height, nil
|
|
}
|
|
|
|
func GetImageFromUrl(url string) (mimeType string, data string, err error) {
|
|
// openai's image_url support base64 encoded image
|
|
if strings.HasPrefix(url, "data:image/jpeg;base64,") {
|
|
return "image/jpeg", strings.TrimPrefix(url, "data:image/jpeg;base64,"), nil
|
|
}
|
|
|
|
isImage, err := IsImageUrl(url)
|
|
if !isImage {
|
|
return
|
|
}
|
|
resp, err := http.Get(url)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
buffer := bytes.NewBuffer(nil)
|
|
_, err = buffer.ReadFrom(resp.Body)
|
|
if err != nil {
|
|
return
|
|
}
|
|
mimeType = resp.Header.Get("Content-Type")
|
|
data = base64.StdEncoding.EncodeToString(buffer.Bytes())
|
|
return
|
|
}
|
|
|
|
var (
|
|
reg = regexp.MustCompile(`data:image/([^;]+);base64,`)
|
|
)
|
|
|
|
var readerPool = sync.Pool{
|
|
New: func() interface{} {
|
|
return &bytes.Reader{}
|
|
},
|
|
}
|
|
|
|
func GetImageSizeFromBase64(encoded string) (width int, height int, err error) {
|
|
decoded, err := base64.StdEncoding.DecodeString(reg.ReplaceAllString(encoded, ""))
|
|
if err != nil {
|
|
return 0, 0, err
|
|
}
|
|
|
|
reader := readerPool.Get().(*bytes.Reader)
|
|
defer readerPool.Put(reader)
|
|
reader.Reset(decoded)
|
|
|
|
img, _, err := image.DecodeConfig(reader)
|
|
if err != nil {
|
|
return 0, 0, err
|
|
}
|
|
|
|
return img.Width, img.Height, nil
|
|
}
|
|
|
|
func GetImageSize(image string) (width int, height int, err error) {
|
|
if strings.HasPrefix(image, "data:image/") {
|
|
return GetImageSizeFromBase64(image)
|
|
}
|
|
return GetImageSizeFromUrl(image)
|
|
}
|