/
nudebox.go
162 lines (150 loc) · 3.97 KB
/
nudebox.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
// Package nudebox provides a client for accessing Nudebox services.
package nudebox
import (
"bytes"
"io"
"mime/multipart"
"net/http"
"net/url"
"strings"
"time"
"github.com/machinebox/sdk-go/boxutil"
"github.com/machinebox/sdk-go/internal/mbhttp"
"github.com/pkg/errors"
)
// Client is an HTTP client that can make requests to the box.
type Client struct {
addr string
// HTTPClient is the http.Client that will be used to
// make requests.
HTTPClient *http.Client
}
// make sure the Client implements boxutil.Box
var _ boxutil.Box = (*Client)(nil)
// New makes a new Client for the box at the specified address.
func New(addr string) *Client {
return &Client{
addr: addr,
HTTPClient: &http.Client{
Timeout: 1 * time.Minute,
},
}
}
// Info gets the details about the box.
func (c *Client) Info() (*boxutil.Info, error) {
var info boxutil.Info
u, err := url.Parse(c.addr + "/info")
if err != nil {
return nil, err
}
if !u.IsAbs() {
return nil, errors.New("box address must be absolute")
}
req, err := http.NewRequest("GET", u.String(), nil)
if err != nil {
return nil, err
}
req.Header.Set("Accept", "application/json; charset=utf-8")
_, err = mbhttp.New("nudebox", c.HTTPClient).DoUnmarshal(req, &info)
if err != nil {
return nil, err
}
return &info, nil
}
// Check gets the nudity probability for the image data provided.
func (c *Client) Check(image io.Reader) (float64, error) {
var buf bytes.Buffer
w := multipart.NewWriter(&buf)
fw, err := w.CreateFormFile("file", "image.dat")
if err != nil {
return 0, err
}
_, err = io.Copy(fw, image)
if err != nil {
return 0, err
}
if err = w.Close(); err != nil {
return 0, err
}
u, err := url.Parse(c.addr + "/nudebox/check")
if err != nil {
return 0, err
}
if !u.IsAbs() {
return 0, errors.New("box address must be absolute")
}
req, err := http.NewRequest("POST", u.String(), &buf)
if err != nil {
return 0, err
}
req.Header.Set("Accept", "application/json; charset=utf-8")
req.Header.Set("Content-Type", w.FormDataContentType())
var checkResponse struct {
Nude float64
}
_, err = mbhttp.New("nudebox", c.HTTPClient).DoUnmarshal(req, &checkResponse)
if err != nil {
return 0, err
}
return checkResponse.Nude, nil
}
// CheckURL gets the nudity probability for the image at the specified URL.
func (c *Client) CheckURL(imageURL *url.URL) (float64, error) {
u, err := url.Parse(c.addr + "/nudebox/check")
if err != nil {
return 0, err
}
if !u.IsAbs() {
return 0, errors.New("box address must be absolute")
}
if !imageURL.IsAbs() {
return 0, errors.New("url must be absolute")
}
form := url.Values{}
form.Set("url", imageURL.String())
req, err := http.NewRequest("POST", u.String(), strings.NewReader(form.Encode()))
if err != nil {
return 0, err
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
req.Header.Set("Accept", "application/json; charset=utf-8")
var checkResponse struct {
Nude float64
}
_, err = mbhttp.New("nudebox", c.HTTPClient).DoUnmarshal(req, &checkResponse)
if err != nil {
return 0, err
}
return checkResponse.Nude, nil
}
// CheckBase64 gets the nudity probability for the Base64 encoded image.
func (c *Client) CheckBase64(data string) (float64, error) {
u, err := url.Parse(c.addr + "/nudebox/check")
if err != nil {
return 0, err
}
if !u.IsAbs() {
return 0, errors.New("box address must be absolute")
}
form := url.Values{}
form.Set("base64", data)
req, err := http.NewRequest("POST", u.String(), strings.NewReader(form.Encode()))
if err != nil {
return 0, err
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
req.Header.Set("Accept", "application/json; charset=utf-8")
var checkResponse struct {
Nude float64
}
_, err = mbhttp.New("nudebox", c.HTTPClient).DoUnmarshal(req, &checkResponse)
if err != nil {
return 0, err
}
return checkResponse.Nude, nil
}
// ErrNudebox represents an error from Nudebox.
type ErrNudebox string
func (e ErrNudebox) Error() string {
return "nudebox: " + string(e)
}