-
Notifications
You must be signed in to change notification settings - Fork 1
/
photo_scan.py
155 lines (131 loc) · 5.03 KB
/
photo_scan.py
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
import requests
import os
class PhotoScanApp:
def __init__(self):
self.plugins = []
def register_plugin(self, plugin):
self.plugins.append(plugin)
def scan_photo(self, photo_path):
analysis_results = []
for plugin in self.plugins:
analysis = plugin.analyze_photo(photo_path)
analysis_results.append(analysis)
return analysis_results
class HiveAI_Plugin:
def check_for_bad_classes(self, json_result, threshold=0.5):
bad_classes = [
"yes_nsfw",
"yes_suggestive",
"yes_female_underwear",
"yes_male_underwear",
"yes_sex_toy",
"yes_female_nudity",
"yes_male_nudity",
"yes_female_swimwear",
"yes_male_shirtless",
"yes_gun_in_hand",
"yes_culinary_knife_in_hand",
"yes_knife_in_hand",
"yes_pills",
"yes_smoking",
"yes_illicit_injectables",
"yes_kkk",
"yes_middle_finger",
"yes_sexual_activity",
"yes_hanging",
"yes_noose",
"yes_realistic_nsfw",
"yes_animated_corpse",
"yes_human_corpse",
"yes_self_harm",
"yes_emaciated_body",
"yes_child_present",
"yes_sexual_intent",
"yes_animal_genitalia_and_human",
"yes_animated_animal_genitalia",
"yes_gambling",
"yes_undressed",
"yes_confederate",
"yes_animated_alcohol",
"yes_alcohol",
"yes_drinking_alcohol",
"yes_religious_icon",
"general_nsfw",
"animated_gun",
"gun_in_hand",
"gun_not_in_hand",
"culinary_knife_in_hand",
"culinary_knife_not_in_hand",
"knife_in_hand",
"knife_not_in_hand",
"a_little_bloody",
"other_blood",
"very_bloody",
"yes_smoking",
"illicit_injectables",
"medical_injectables",
"hanging",
"animated_corpse",
"human_corpse",
"animal_genitalia_and_human",
"animal_genitalia_only",
"animated_animal_genitalia",
"animated_alcohol",
"general_suggestive",
]
output = json_result.get("status", [])[0].get("response", {}).get("output", [])
harmful_classes = []
for item in output:
classes = item.get("classes", [])
for class_info in classes:
class_name = class_info.get("class", "")
score = class_info.get("score", 0)
if class_name in bad_classes and score >= threshold:
harmful_classes.append(class_name.strip("yes_"))
if harmful_classes:
return True, ", ".join(harmful_classes)
else:
return False, None
def analyze_photo(self,url):
payload={"url": url}
headers = {
"accept": "application/json",
"authorization": "token {}".format(os.environ["HIVEAI_KEY"] )
}
response = requests.post( "https://api.thehive.ai/api/v2/task/sync", headers=headers, data=payload)
response=response.json()
harmfull_content, bad_class = self.check_for_bad_classes(response)
if harmfull_content:
return {"plugin": "HiveAI", "result": f"Harmfull content found: {bad_class}"}
else:
return {"plugin": "HiveAI", "result": "Harmfull content not found"}
class PhotoDNA_Plugin:
def analyze_photo(self,image_url):
url = "https://api.microsoftmoderator.com/photodna/v1.0/Match?enhance=false"
headers = {
"Content-Type": "application/json; charset=utf-8",
"Ocp-Apim-Subscription-Key": os.environ["PHOTODNA_KEY"]
}
data = {
"DataRepresentation": "URL",
"Value": image_url
}
response = requests.post(url, headers=headers, json=data)
scan_result=response.json()
if scan_result["IsMatch"]:
{"plugin": "PhotoDNA", "result": f"PhotoDNA match"}
else:
return {"plugin": "PhotoDNA", "result": f"No PhotoDNA match"}
def initialize_photo_scan_app():
photo_scan_app = PhotoScanApp()
face_detection_plugin = HiveAI_Plugin()
object_recognition_plugin = PhotoDNA_Plugin()
photo_scan_app.register_plugin(face_detection_plugin)
photo_scan_app.register_plugin(object_recognition_plugin)
return photo_scan_app
if __name__=="__main__":
photo_path ="https://s3.amazonaws.com/docs.thehive.ai/client_demo/moderation_image.png"
photo_scan_app=initialize_photo_scan_app()
analysis_results = photo_scan_app.scan_photo(photo_path)
for analysis in analysis_results:
print(analysis)