Create app.py

app.py

@@ -0,0 +1,257 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import cv2
+import uuid
+import os
+import requests
+import random
+import string
+import json
+import shutil
+
+app = FastAPI()
+
+def check_nsfw(img_url):
+    session_hash = ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(7))
+    data = {
+        'data': [
+            {'path': img_url},
+            "chen-convnext",
+            0.5,
+            True,
+            True
+            ],
+        'session_hash': session_hash,
+        'fn_index': 0,
+        'trigger_id': 12
+        }
+    r = requests.post('https://yoinked-da-nsfw-checker.hf.space/queue/join', json=data)
+
+    r = requests.get(f'https://yoinked-da-nsfw-checker.hf.space/queue/data?session_hash={session_hash}', stream=True)
+    buffer = ""  # Buffer to accumulate the chunks
+
+
+    for content in r.iter_content(100):
+        # Decode the byte content to a string
+        buffer += content.decode('utf-8')
+
+    return json.loads(buffer.split('data:')[len(buffer.split('data:'))-2])["output"]["data"][0]['label']
+    
+
+def check_nsfw2(img_url):
+    session_hash = ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(7))
+    data = {
+        'data': [
+            {'path': img_url}
+            ],
+        'session_hash': session_hash,
+        'fn_index': 0,
+        'trigger_id': 9
+        }
+    r = requests.post('https://jamescookjr90-falconsai-nsfw-image-detection.hf.space/queue/join', json=data)
+
+    r = requests.get(f'https://jamescookjr90-falconsai-nsfw-image-detection.hf.space/queue/data?session_hash={session_hash}', stream=True)
+    buffer = ""  # Buffer to accumulate the chunks
+
+
+    for content in r.iter_content(100):
+        # Decode the byte content to a string
+        buffer += content.decode('utf-8')
+    return json.loads(buffer.split('data:')[len(buffer.split('data:'))-2])["output"]["data"][0]['label']
+
+def check_nsfw3(img_url):
+    data = {
+        'data': [
+            {'path': img_url}
+            ]
+
+        }
+    r = requests.post('https://zanderlewis-xl-nsfw-detection.hf.space/call/predict',json=data)
+    json_data = r.json()
+    event_id = json_data['event_id']
+    r = requests.get(f'https://zanderlewis-xl-nsfw-detection.hf.space/call/predict/{event_id}', stream=True)
+    event_stream = ''
+    for chunk in r.iter_content(100):
+        event_stream += chunk.decode('utf-8')
+    return ast.literal_eval(event_stream.split('data:')[-1])[0]['label']
+
+def get_replica_code(url):
+    try:
+        r = requests.get(url)
+
+        return r.text.split('replicas/')[1].split('"};')[0]
+    except:
+        return None
+def check_nsfw4(img_url):
+    code = get_replica_code('https://error466-falconsai-nsfw-image-detection.hf.space')
+    session_hash = ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(7))
+    data = {
+        'data': [
+            {'path': img_url}
+            ],
+        'session_hash': session_hash,
+        'fn_index': 0,
+        'trigger_id': 58
+        }
+    r = requests.post(f'https://error466-falconsai-nsfw-image-detection.hf.space/--replicas/{code}/queue/join', json=data)
+
+    r = requests.get(f'https://error466-falconsai-nsfw-image-detection.hf.space/--replicas/{code}/queue/data?session_hash={session_hash}', stream=True)
+    buffer = ""  # Buffer to accumulate the chunks
+
+
+    for content in r.iter_content(100):
+        # Decode the byte content to a string
+        buffer += content.decode('utf-8')
+    return json.loads(buffer.split('data:')[-1])["output"]["data"][0]['label']
+
+def check_nsfw5(img_url):
+    session_hash = ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(7))
+    data = {
+        'data': [
+            {'path': img_url}
+            ],
+        'session_hash': session_hash,
+        'fn_index': 0,
+        'trigger_id': 9
+        }
+    r = requests.post('https://phelpsgg-falconsai-nsfw-image-detection.hf.space/queue/join', json=data)
+
+    r = requests.get(f'https://phelpsgg-falconsai-nsfw-image-detection.hf.space/queue/data?session_hash={session_hash}', stream=True)
+    buffer = ""  # Buffer to accumulate the chunks
+
+
+    for content in r.iter_content(100):
+        # Decode the byte content to a string
+        buffer += content.decode('utf-8')
+    return json.loads(buffer.split('data:')[-1])["output"]["data"][0]['label']
+    
+def check_nsfw_final(img_url):
+    result = None
+    print(img_url)
+    try:
+        check = check_nsfw2(img_url)
+        if check == 'nsfw':
+            result = True
+        else:
+            result = False
+    except Exception as e:
+        print(e)
+
+    if result is None:
+        try:
+            check = check_nsfw3(img_url)
+            if check == 'nsfw':
+                result = True
+            else:
+                result = False
+        except Exception as e:
+            print(e)
+    
+    if result is None:
+        try:
+            check = check_nsfw4(img_url)
+            if check == 'nsfw':
+                result = True
+            else:
+                result = False
+        except:
+            pass
+
+    if result is None:
+        try:
+            check = check_nsfw5(img_url)
+            if check == 'nsfw':
+                result = True
+            else:
+                result = False
+        except Exception as e:
+            print(e)
+
+    return result
+
+# Frame extraction
+def extract_frames(video_path, num_frames, temp_dir):
+    vidcap = cv2.VideoCapture(video_path)
+    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    frames = []
+    
+    for i in range(num_frames):
+        frame_number = int(i * total_frames / num_frames) 
+        vidcap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
+        success, image = vidcap.read()
+        if success:
+            frame_filename = os.path.join(temp_dir, f"{uuid.uuid4().hex}.jpg")
+            cv2.imwrite(frame_filename, image)
+            frames.append(frame_filename)
+        else:
+            break
+
+    vidcap.release()
+    return frames
+
+# Video processing
+def process_video(video_path, num_frames):
+    temp_dir = os.path.join("/tmp", f"frames_{uuid.uuid4().hex}")
+    os.makedirs(temp_dir, exist_ok=True)
+
+    frames = extract_frames(video_path, num_frames, temp_dir)
+    
+    nsfw_count = 0
+    total_frames = len(frames)
+    frame_results = []
+
+    for frame_path in frames:
+        img_url = f"file://{frame_path}"
+        nsfw_detected = check_nsfw_final(img_url)
+        frame_results.append({
+            "frame_path": frame_path,
+            "nsfw_detected": nsfw_detected
+        })
+        if nsfw_detected:
+            nsfw_count += 1
+    
+    # Cleanup
+    shutil.rmtree(temp_dir, ignore_errors=True)
+
+    result = {
+        "nsfw_count": nsfw_count,
+        "total_frames": total_frames,
+        "frames": frame_results
+    }
+    
+    return result
+
+# Request/Response models
+class VideoRequest(BaseModel):
+    video_url: str
+    num_frames: int = 10
+
+class VideoResponse(BaseModel):
+    nsfw_count: int
+    total_frames: int
+    frames: list
+
+# API Endpoints
+
+@app.post("/process_video", response_model=VideoResponse)
+async def process_video_endpoint(request: VideoRequest):
+    # Download video
+    try:
+        video_path = os.path.join("/tmp", f"{uuid.uuid4().hex}.mp4")
+        with requests.get(request.video_url, stream=True) as r:
+            r.raise_for_status()
+            with open(video_path, "wb") as f:
+                for chunk in r.iter_content(chunk_size=8192):
+                    f.write(chunk)
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error downloading video: {e}")
+    
+    try:
+        # Process video
+        result = process_video(video_path, request.num_frames)
+    finally:
+        # Cleanup video file
+        if os.path.exists(video_path):
+            os.remove(video_path)
+    
+    return result