diff --git a/automod/visual/doc.go b/automod/visual/doc.go new file mode 100644 index 00000000..2798b2ec --- /dev/null +++ b/automod/visual/doc.go @@ -0,0 +1,3 @@ +// automod helpers for visual content (image blobs) +package visual + diff --git a/automod/visual/hiveai.go b/automod/visual/hiveai.go new file mode 100644 index 00000000..21594f60 --- /dev/null +++ b/automod/visual/hiveai.go @@ -0,0 +1,147 @@ +package visual + +import ( + "log/slog" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "mime/multipart" + "net/http" + + lexutil "github.com/bluesky-social/indigo/lex/util" + "github.com/bluesky-social/indigo/util" + + "github.com/carlmjohnson/versioninfo" +) + +type HiveAILabeler struct { + Client http.Client + ApiToken string +} + +// schema: https://docs.thehive.ai/reference/classification +type HiveAIResp struct { + Status []HiveAIResp_Status `json:"status"` +} + +type HiveAIResp_Status struct { + Response HiveAIResp_Response `json:"response"` +} + +type HiveAIResp_Response struct { + Output []HiveAIResp_Out `json:"output"` +} + +type HiveAIResp_Out struct { + Time float64 `json:"time"` + Classes []HiveAIResp_Class `json:"classes"` +} + +type HiveAIResp_Class struct { + Class string `json:"class"` + Score float64 `json:"score"` +} + +func NewHiveAILabeler(token string) HiveAILabeler { + return HiveAILabeler{ + Client: *util.RobustHTTPClient(), + ApiToken: token, + } +} + +func (resp *HiveAIResp) SummarizeLabels() []string { + var labels []string + + for _, status := range resp.Status { + for _, out := range status.Response.Output { + for _, cls := range out.Classes { + // TODO(bnewbold): lots more upstream tags could be included here. + // for example, "sexy" for not nude but still explicit/suggestive, + // or lolicon (animated, not nude, "sugggestive" + + // sexual: https://docs.thehive.ai/docs/sexual-content + // note: won't apply "nude" if "porn" already applied + if cls.Class == "yes_sexual_activity" && cls.Score >= 0.90 { + // NOTE: will include "hentai" + labels = append(labels, "porn") + } else if cls.Class == "animal_genitalia_and_human" && cls.Score >= 0.90 { + labels = append(labels, "porn") + } else if cls.Class == "yes_male_nudity" && cls.Score >= 0.90 { + labels = append(labels, "nude") + } else if cls.Class == "yes_female_nudity" && cls.Score >= 0.90 { + labels = append(labels, "nude") + } + + // gore and violence: https://docs.thehive.ai/docs/class-descriptions-violence-gore + if cls.Class == "very_bloody" && cls.Score >= 0.90 { + labels = append(labels, "gore") + } + if cls.Class == "human_corpse" && cls.Score >= 0.90 { + labels = append(labels, "corpse") + } + if cls.Class == "yes_self_harm" && cls.Score >= 0.90 { + labels = append(labels, "self-harm") + } + } + } + } + + return labels +} + +func (hal *HiveAILabeler) LabelBlob(ctx context.Context, blob lexutil.LexBlob, blobBytes []byte) ([]string, error) { + + slog.Info("sending blob to thehive.ai", "cid", blob.Ref, "mimetype", blob.MimeType, "size", len(blobBytes)) + + // generic HTTP form file upload, then parse the response JSON + body := &bytes.Buffer{} + writer := multipart.NewWriter(body) + part, err := writer.CreateFormFile("media", blob.Ref.String()) + if err != nil { + return nil, err + } + _, err = part.Write(blobBytes) + if err != nil { + return nil, err + } + err = writer.Close() + if err != nil { + return nil, err + } + + req, err := http.NewRequest("POST", "https://api.thehive.ai/api/v2/task/sync", body) + if err != nil { + return nil, err + } + + req.Header.Set("Authorization", fmt.Sprintf("Token %s", hal.ApiToken)) + req.Header.Add("Content-Type", writer.FormDataContentType()) + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "labelmaker/"+versioninfo.Short()) + + res, err := hal.Client.Do(req) + if err != nil { + return nil, fmt.Errorf("HiveAI request failed: %v", err) + } + defer res.Body.Close() + if res.StatusCode != 200 { + return nil, fmt.Errorf("HiveAI request failed statusCode=%d", res.StatusCode) + } + + respBytes, err := io.ReadAll(res.Body) + if err != nil { + return nil, fmt.Errorf("failed to read HiveAI resp body: %v", err) + } + + slog.Debug("HiveAI raw result", "cid", blob.Ref, "body", string(respBytes)) + + var respObj HiveAIResp + if err := json.Unmarshal(respBytes, &respObj); err != nil { + return nil, fmt.Errorf("failed to parse HiveAI resp JSON: %v", err) + } + respJson, _ := json.Marshal(respObj.Status[0].Response.Output[0]) + slog.Info("HiveAI result", "cid", blob.Ref, "json", string(respJson)) + return respObj.SummarizeLabels(), nil +} diff --git a/automod/visual/hiveai_test.go b/automod/visual/hiveai_test.go new file mode 100644 index 00000000..3e79aa26 --- /dev/null +++ b/automod/visual/hiveai_test.go @@ -0,0 +1,42 @@ +package visual + +import ( + "encoding/json" + "io" + "os" + "reflect" + "testing" +) + +func TestHiveParse(t *testing.T) { + file, err := os.Open("testdata/hiveai_resp_example.json") + if err != nil { + t.Fatal(err) + } + + respBytes, err := io.ReadAll(file) + if err != nil { + t.Fatal(err) + } + + var respObj HiveAIResp + if err := json.Unmarshal(respBytes, &respObj); err != nil { + t.Fatal(err) + } + + classes := respObj.Status[0].Response.Output[0].Classes + if len(classes) <= 10 { + t.Fatal("didn't get expected class count") + } + for _, c := range classes { + if c.Class == "" || c.Score == 0.0 { + t.Fatal("got null/empty class in resp") + } + } + + labels := respObj.SummarizeLabels() + expected := []string{"porn"} + if !reflect.DeepEqual(labels, expected) { + t.Fatal("didn't summarize to expected labels") + } +} diff --git a/automod/visual/testdata/hiveai_resp_example.json b/automod/visual/testdata/hiveai_resp_example.json new file mode 100644 index 00000000..2a378604 --- /dev/null +++ b/automod/visual/testdata/hiveai_resp_example.json @@ -0,0 +1,401 @@ +{ + "id": "02122580-c37f-11ed-81d2-000000000000", + "code": 200, + "project_id": 12345, + "user_id": 12345, + "created_on": "2023-03-15T22:16:18.408Z", + "status": [ + { + "status": { + "code": "0", + "message": "SUCCESS" + }, + "response": { + "input": { + "id": "02122580-c37f-11ed-81d2-000000000000", + "charge": 0.003, + "model": "mod55_dense", + "model_version": 1, + "model_type": "CATEGORIZATION", + "created_on": "2023-03-15T22:16:18.136Z", + "media": { + "url": null, + "filename": "bafkreiam7k6mvkyuoybq4ynhljvj5xa75sdbhjbolzjf5j2udx7vj5gnsy", + "type": "PHOTO", + "mime_type": "jpeg", + "mimetype": "image/jpeg", + "width": 800, + "height": 800, + "num_frames": 1, + "duration": 0 + }, + "user_id": 12345, + "project_id": 12345, + "config_version": 1, + "config_tag": "default" + }, + "output": [ + { + "time": 0, + "classes": [ + { + "class": "general_not_nsfw_not_suggestive", + "score": 0.9998097218132356 + }, + { + "class": "general_nsfw", + "score": 8.857344804177162e-05 + }, + { + "class": "general_suggestive", + "score": 0.00010170473872266839 + }, + { + "class": "no_female_underwear", + "score": 0.9999923079040384 + }, + { + "class": "yes_female_underwear", + "score": 7.692095961599136e-06 + }, + { + "class": "no_male_underwear", + "score": 0.9999984904867634 + }, + { + "class": "yes_male_underwear", + "score": 1.5095132367094679e-06 + }, + { + "class": "no_sex_toy", + "score": 0.9999970970762551 + }, + { + "class": "yes_sex_toy", + "score": 2.9029237450490604e-06 + }, + { + "class": "no_female_nudity", + "score": 0.9999739028909301 + }, + { + "class": "yes_female_nudity", + "score": 2.60971090699536e-05 + }, + { + "class": "no_male_nudity", + "score": 0.9999711373083747 + }, + { + "class": "yes_male_nudity", + "score": 2.8862691625255323e-05 + }, + { + "class": "no_female_swimwear", + "score": 0.9999917609899659 + }, + { + "class": "yes_female_swimwear", + "score": 8.239010034025379e-06 + }, + { + "class": "no_male_shirtless", + "score": 0.9999583350744331 + }, + { + "class": "yes_male_shirtless", + "score": 4.166492556688088e-05 + }, + { + "class": "no_text", + "score": 0.9958378716447616 + }, + { + "class": "text", + "score": 0.0041621283552384265 + }, + { + "class": "animated", + "score": 0.46755478950048235 + }, + { + "class": "hybrid", + "score": 0.0011440363434524984 + }, + { + "class": "natural", + "score": 0.5313011741560651 + }, + { + "class": "animated_gun", + "score": 2.0713000782979496e-05 + }, + { + "class": "gun_in_hand", + "score": 1.5844730446534659e-06 + }, + { + "class": "gun_not_in_hand", + "score": 1.0338973818006654e-06 + }, + { + "class": "no_gun", + "score": 0.9999766686287906 + }, + { + "class": "culinary_knife_in_hand", + "score": 3.8063500083369785e-06 + }, + { + "class": "culinary_knife_not_in_hand", + "score": 7.94057948996249e-07 + }, + { + "class": "knife_in_hand", + "score": 4.5578955723278505e-07 + }, + { + "class": "knife_not_in_hand", + "score": 3.842124714748908e-07 + }, + { + "class": "no_knife", + "score": 0.999994559590014 + }, + { + "class": "a_little_bloody", + "score": 2.1317745626539786e-07 + }, + { + "class": "no_blood", + "score": 0.9999793341236429 + }, + { + "class": "other_blood", + "score": 2.0322054269591763e-05 + }, + { + "class": "very_bloody", + "score": 1.306446309561673e-07 + }, + { + "class": "no_pills", + "score": 0.9999989592376954 + }, + { + "class": "yes_pills", + "score": 1.0407623044588633e-06 + }, + { + "class": "no_smoking", + "score": 0.9999939101969173 + }, + { + "class": "yes_smoking", + "score": 6.089803082758281e-06 + }, + { + "class": "illicit_injectables", + "score": 6.925695592003094e-07 + }, + { + "class": "medical_injectables", + "score": 8.587808234452378e-07 + }, + { + "class": "no_injectables", + "score": 0.9999984486496174 + }, + { + "class": "no_nazi", + "score": 0.9999987449628097 + }, + { + "class": "yes_nazi", + "score": 1.2550371902234279e-06 + }, + { + "class": "no_kkk", + "score": 0.999999762417549 + }, + { + "class": "yes_kkk", + "score": 2.3758245111050425e-07 + }, + { + "class": "no_middle_finger", + "score": 0.9999881515231847 + }, + { + "class": "yes_middle_finger", + "score": 1.184847681536747e-05 + }, + { + "class": "no_terrorist", + "score": 0.9999998870793229 + }, + { + "class": "yes_terrorist", + "score": 1.1292067715380635e-07 + }, + { + "class": "no_overlay_text", + "score": 0.9996453363440359 + }, + { + "class": "yes_overlay_text", + "score": 0.0003546636559640924 + }, + { + "class": "no_sexual_activity", + "score": 0.9999563580374798 + }, + { + "class": "yes_sexual_activity", + "score": 0.99, + "realScore": 4.364196252012032e-05 + }, + { + "class": "hanging", + "score": 3.6435135762510905e-07 + }, + { + "class": "no_hanging_no_noose", + "score": 0.9999980779196416 + }, + { + "class": "noose", + "score": 1.5577290007796094e-06 + }, + { + "class": "no_realistic_nsfw", + "score": 0.9999944341007805 + }, + { + "class": "yes_realistic_nsfw", + "score": 5.565899219571182e-06 + }, + { + "class": "animated_corpse", + "score": 5.276802046755426e-07 + }, + { + "class": "human_corpse", + "score": 2.5449360984211012e-08 + }, + { + "class": "no_corpse", + "score": 0.9999994468704343 + }, + { + "class": "no_self_harm", + "score": 0.9999994515625507 + }, + { + "class": "yes_self_harm", + "score": 5.484374493605692e-07 + }, + { + "class": "no_drawing", + "score": 0.9978276028816608 + }, + { + "class": "yes_drawing", + "score": 0.0021723971183392485 + }, + { + "class": "no_emaciated_body", + "score": 0.9999998146500432 + }, + { + "class": "yes_emaciated_body", + "score": 1.853499568724518e-07 + }, + { + "class": "no_child_present", + "score": 0.9999970498515446 + }, + { + "class": "yes_child_present", + "score": 2.950148455380443e-06 + }, + { + "class": "no_sexual_intent", + "score": 0.9999963861546292 + }, + { + "class": "yes_sexual_intent", + "score": 3.613845370766111e-06 + }, + { + "class": "animal_genitalia_and_human", + "score": 2.255472023465222e-08 + }, + { + "class": "animal_genitalia_only", + "score": 4.6783185199931176e-07 + }, + { + "class": "animated_animal_genitalia", + "score": 6.707857419436447e-07 + }, + { + "class": "no_animal_genitalia", + "score": 0.9999988388276858 + }, + { + "class": "no_gambling", + "score": 0.9999960939687145 + }, + { + "class": "yes_gambling", + "score": 3.906031285604864e-06 + }, + { + "class": "no_undressed", + "score": 0.99999923356218 + }, + { + "class": "yes_undressed", + "score": 7.664378199789045e-07 + }, + { + "class": "no_confederate", + "score": 0.9999925456900376 + }, + { + "class": "yes_confederate", + "score": 7.454309962453175e-06 + }, + { + "class": "animated_alcohol", + "score": 1.8109949948066074e-06 + }, + { + "class": "no_alcohol", + "score": 0.9999916620957963 + }, + { + "class": "yes_alcohol", + "score": 5.88781463445443e-06 + }, + { + "class": "yes_drinking_alcohol", + "score": 6.390945746578106e-07 + }, + { + "class": "no_religious_icon", + "score": 0.9999862158580689 + }, + { + "class": "yes_religious_icon", + "score": 1.3784141931119298e-05 + } + ] + } + ] + } + } + ], + "from_cache": false +}