sdafd commited on
Commit
6f695bf
·
verified ·
1 Parent(s): 7cc5ea8

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +257 -0
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