Spaces:
Sleeping
Sleeping
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,257 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from fastapi import FastAPI, HTTPException
|
2 |
+
from pydantic import BaseModel
|
3 |
+
import cv2
|
4 |
+
import uuid
|
5 |
+
import os
|
6 |
+
import requests
|
7 |
+
import random
|
8 |
+
import string
|
9 |
+
import json
|
10 |
+
import shutil
|
11 |
+
|
12 |
+
app = FastAPI()
|
13 |
+
|
14 |
+
def check_nsfw(img_url):
|
15 |
+
session_hash = ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(7))
|
16 |
+
data = {
|
17 |
+
'data': [
|
18 |
+
{'path': img_url},
|
19 |
+
"chen-convnext",
|
20 |
+
0.5,
|
21 |
+
True,
|
22 |
+
True
|
23 |
+
],
|
24 |
+
'session_hash': session_hash,
|
25 |
+
'fn_index': 0,
|
26 |
+
'trigger_id': 12
|
27 |
+
}
|
28 |
+
r = requests.post('https://yoinked-da-nsfw-checker.hf.space/queue/join', json=data)
|
29 |
+
|
30 |
+
r = requests.get(f'https://yoinked-da-nsfw-checker.hf.space/queue/data?session_hash={session_hash}', stream=True)
|
31 |
+
buffer = "" # Buffer to accumulate the chunks
|
32 |
+
|
33 |
+
|
34 |
+
for content in r.iter_content(100):
|
35 |
+
# Decode the byte content to a string
|
36 |
+
buffer += content.decode('utf-8')
|
37 |
+
|
38 |
+
return json.loads(buffer.split('data:')[len(buffer.split('data:'))-2])["output"]["data"][0]['label']
|
39 |
+
|
40 |
+
|
41 |
+
def check_nsfw2(img_url):
|
42 |
+
session_hash = ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(7))
|
43 |
+
data = {
|
44 |
+
'data': [
|
45 |
+
{'path': img_url}
|
46 |
+
],
|
47 |
+
'session_hash': session_hash,
|
48 |
+
'fn_index': 0,
|
49 |
+
'trigger_id': 9
|
50 |
+
}
|
51 |
+
r = requests.post('https://jamescookjr90-falconsai-nsfw-image-detection.hf.space/queue/join', json=data)
|
52 |
+
|
53 |
+
r = requests.get(f'https://jamescookjr90-falconsai-nsfw-image-detection.hf.space/queue/data?session_hash={session_hash}', stream=True)
|
54 |
+
buffer = "" # Buffer to accumulate the chunks
|
55 |
+
|
56 |
+
|
57 |
+
for content in r.iter_content(100):
|
58 |
+
# Decode the byte content to a string
|
59 |
+
buffer += content.decode('utf-8')
|
60 |
+
return json.loads(buffer.split('data:')[len(buffer.split('data:'))-2])["output"]["data"][0]['label']
|
61 |
+
|
62 |
+
def check_nsfw3(img_url):
|
63 |
+
data = {
|
64 |
+
'data': [
|
65 |
+
{'path': img_url}
|
66 |
+
]
|
67 |
+
|
68 |
+
}
|
69 |
+
r = requests.post('https://zanderlewis-xl-nsfw-detection.hf.space/call/predict',json=data)
|
70 |
+
json_data = r.json()
|
71 |
+
event_id = json_data['event_id']
|
72 |
+
r = requests.get(f'https://zanderlewis-xl-nsfw-detection.hf.space/call/predict/{event_id}', stream=True)
|
73 |
+
event_stream = ''
|
74 |
+
for chunk in r.iter_content(100):
|
75 |
+
event_stream += chunk.decode('utf-8')
|
76 |
+
return ast.literal_eval(event_stream.split('data:')[-1])[0]['label']
|
77 |
+
|
78 |
+
def get_replica_code(url):
|
79 |
+
try:
|
80 |
+
r = requests.get(url)
|
81 |
+
|
82 |
+
return r.text.split('replicas/')[1].split('"};')[0]
|
83 |
+
except:
|
84 |
+
return None
|
85 |
+
def check_nsfw4(img_url):
|
86 |
+
code = get_replica_code('https://error466-falconsai-nsfw-image-detection.hf.space')
|
87 |
+
session_hash = ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(7))
|
88 |
+
data = {
|
89 |
+
'data': [
|
90 |
+
{'path': img_url}
|
91 |
+
],
|
92 |
+
'session_hash': session_hash,
|
93 |
+
'fn_index': 0,
|
94 |
+
'trigger_id': 58
|
95 |
+
}
|
96 |
+
r = requests.post(f'https://error466-falconsai-nsfw-image-detection.hf.space/--replicas/{code}/queue/join', json=data)
|
97 |
+
|
98 |
+
r = requests.get(f'https://error466-falconsai-nsfw-image-detection.hf.space/--replicas/{code}/queue/data?session_hash={session_hash}', stream=True)
|
99 |
+
buffer = "" # Buffer to accumulate the chunks
|
100 |
+
|
101 |
+
|
102 |
+
for content in r.iter_content(100):
|
103 |
+
# Decode the byte content to a string
|
104 |
+
buffer += content.decode('utf-8')
|
105 |
+
return json.loads(buffer.split('data:')[-1])["output"]["data"][0]['label']
|
106 |
+
|
107 |
+
def check_nsfw5(img_url):
|
108 |
+
session_hash = ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(7))
|
109 |
+
data = {
|
110 |
+
'data': [
|
111 |
+
{'path': img_url}
|
112 |
+
],
|
113 |
+
'session_hash': session_hash,
|
114 |
+
'fn_index': 0,
|
115 |
+
'trigger_id': 9
|
116 |
+
}
|
117 |
+
r = requests.post('https://phelpsgg-falconsai-nsfw-image-detection.hf.space/queue/join', json=data)
|
118 |
+
|
119 |
+
r = requests.get(f'https://phelpsgg-falconsai-nsfw-image-detection.hf.space/queue/data?session_hash={session_hash}', stream=True)
|
120 |
+
buffer = "" # Buffer to accumulate the chunks
|
121 |
+
|
122 |
+
|
123 |
+
for content in r.iter_content(100):
|
124 |
+
# Decode the byte content to a string
|
125 |
+
buffer += content.decode('utf-8')
|
126 |
+
return json.loads(buffer.split('data:')[-1])["output"]["data"][0]['label']
|
127 |
+
|
128 |
+
def check_nsfw_final(img_url):
|
129 |
+
result = None
|
130 |
+
print(img_url)
|
131 |
+
try:
|
132 |
+
check = check_nsfw2(img_url)
|
133 |
+
if check == 'nsfw':
|
134 |
+
result = True
|
135 |
+
else:
|
136 |
+
result = False
|
137 |
+
except Exception as e:
|
138 |
+
print(e)
|
139 |
+
|
140 |
+
if result is None:
|
141 |
+
try:
|
142 |
+
check = check_nsfw3(img_url)
|
143 |
+
if check == 'nsfw':
|
144 |
+
result = True
|
145 |
+
else:
|
146 |
+
result = False
|
147 |
+
except Exception as e:
|
148 |
+
print(e)
|
149 |
+
|
150 |
+
if result is None:
|
151 |
+
try:
|
152 |
+
check = check_nsfw4(img_url)
|
153 |
+
if check == 'nsfw':
|
154 |
+
result = True
|
155 |
+
else:
|
156 |
+
result = False
|
157 |
+
except:
|
158 |
+
pass
|
159 |
+
|
160 |
+
if result is None:
|
161 |
+
try:
|
162 |
+
check = check_nsfw5(img_url)
|
163 |
+
if check == 'nsfw':
|
164 |
+
result = True
|
165 |
+
else:
|
166 |
+
result = False
|
167 |
+
except Exception as e:
|
168 |
+
print(e)
|
169 |
+
|
170 |
+
return result
|
171 |
+
|
172 |
+
# Frame extraction
|
173 |
+
def extract_frames(video_path, num_frames, temp_dir):
|
174 |
+
vidcap = cv2.VideoCapture(video_path)
|
175 |
+
total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
|
176 |
+
frames = []
|
177 |
+
|
178 |
+
for i in range(num_frames):
|
179 |
+
frame_number = int(i * total_frames / num_frames)
|
180 |
+
vidcap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
|
181 |
+
success, image = vidcap.read()
|
182 |
+
if success:
|
183 |
+
frame_filename = os.path.join(temp_dir, f"{uuid.uuid4().hex}.jpg")
|
184 |
+
cv2.imwrite(frame_filename, image)
|
185 |
+
frames.append(frame_filename)
|
186 |
+
else:
|
187 |
+
break
|
188 |
+
|
189 |
+
vidcap.release()
|
190 |
+
return frames
|
191 |
+
|
192 |
+
# Video processing
|
193 |
+
def process_video(video_path, num_frames):
|
194 |
+
temp_dir = os.path.join("/tmp", f"frames_{uuid.uuid4().hex}")
|
195 |
+
os.makedirs(temp_dir, exist_ok=True)
|
196 |
+
|
197 |
+
frames = extract_frames(video_path, num_frames, temp_dir)
|
198 |
+
|
199 |
+
nsfw_count = 0
|
200 |
+
total_frames = len(frames)
|
201 |
+
frame_results = []
|
202 |
+
|
203 |
+
for frame_path in frames:
|
204 |
+
img_url = f"file://{frame_path}"
|
205 |
+
nsfw_detected = check_nsfw_final(img_url)
|
206 |
+
frame_results.append({
|
207 |
+
"frame_path": frame_path,
|
208 |
+
"nsfw_detected": nsfw_detected
|
209 |
+
})
|
210 |
+
if nsfw_detected:
|
211 |
+
nsfw_count += 1
|
212 |
+
|
213 |
+
# Cleanup
|
214 |
+
shutil.rmtree(temp_dir, ignore_errors=True)
|
215 |
+
|
216 |
+
result = {
|
217 |
+
"nsfw_count": nsfw_count,
|
218 |
+
"total_frames": total_frames,
|
219 |
+
"frames": frame_results
|
220 |
+
}
|
221 |
+
|
222 |
+
return result
|
223 |
+
|
224 |
+
# Request/Response models
|
225 |
+
class VideoRequest(BaseModel):
|
226 |
+
video_url: str
|
227 |
+
num_frames: int = 10
|
228 |
+
|
229 |
+
class VideoResponse(BaseModel):
|
230 |
+
nsfw_count: int
|
231 |
+
total_frames: int
|
232 |
+
frames: list
|
233 |
+
|
234 |
+
# API Endpoints
|
235 |
+
|
236 |
+
@app.post("/process_video", response_model=VideoResponse)
|
237 |
+
async def process_video_endpoint(request: VideoRequest):
|
238 |
+
# Download video
|
239 |
+
try:
|
240 |
+
video_path = os.path.join("/tmp", f"{uuid.uuid4().hex}.mp4")
|
241 |
+
with requests.get(request.video_url, stream=True) as r:
|
242 |
+
r.raise_for_status()
|
243 |
+
with open(video_path, "wb") as f:
|
244 |
+
for chunk in r.iter_content(chunk_size=8192):
|
245 |
+
f.write(chunk)
|
246 |
+
except Exception as e:
|
247 |
+
raise HTTPException(status_code=400, detail=f"Error downloading video: {e}")
|
248 |
+
|
249 |
+
try:
|
250 |
+
# Process video
|
251 |
+
result = process_video(video_path, request.num_frames)
|
252 |
+
finally:
|
253 |
+
# Cleanup video file
|
254 |
+
if os.path.exists(video_path):
|
255 |
+
os.remove(video_path)
|
256 |
+
|
257 |
+
return result
|