Spaces:
Runtime error
Runtime error
File size: 2,990 Bytes
918f252 b98fea5 d9e40e8 b98fea5 d9e40e8 b98fea5 d9e40e8 b98fea5 d9e40e8 b98fea5 d9e40e8 b98fea5 918f252 d9e40e8 918f252 d9e40e8 918f252 d9e40e8 918f252 b98fea5 d9e40e8 b98fea5 d9e40e8 b98fea5 d9e40e8 918f252 d9e40e8 918f252 b98fea5 918f252 d9e40e8 918f252 b98fea5 918f252 |
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 |
import os
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
import torch
from flask_cors import CORS
from flask import Flask, request, json
from PIL import Image
import requests
from io import BytesIO
from bs4 import BeautifulSoup
from urllib.parse import urljoin
os.environ["CUDA_VISIBLE_DEVICES"] = ""
app = Flask(__name__)
cors = CORS(app)
# Define the model and feature extractor globally
model = AutoModelForImageClassification.from_pretrained('carbon225/vit-base-patch16-224-hentai')
feature_extractor = AutoFeatureExtractor.from_pretrained('carbon225/vit-base-patch16-224-hentai')
def predict(response):
try:
# Open and preprocess the image
image = Image.open(BytesIO(response.content))
image = image.resize((128, 128))
# Extract features using the pre-trained feature extractor
encoding = feature_extractor(images=image.convert("RGB"), return_tensors="pt")
# Make a prediction using the pre-trained model
with torch.no_grad():
outputs = model(**encoding)
logits = outputs.logits
# Get the predicted class index and label
predicted_class_idx = logits.argmax(-1).item()
predicted_class_label = model.config.id2label[predicted_class_idx]
return predicted_class_label
except Exception as e:
print(f"Error in predicting image: {str(e)}")
return None
@app.route("/", methods=["GET"])
def default():
return json.dumps({"Server": "Working"})
@app.route("/extractimages", methods=["GET"])
def extract_images():
try:
src = request.args.get("src")
response = requests.get(src)
soup = BeautifulSoup(response.content, 'html.parser')
img_tags = soup.select('div img')
for img_tag in img_tags:
img_url = urljoin(src, img_tag['src'])
response = requests.get(img_url)
response.raise_for_status()
predicted_class_label = predict(response)
if predicted_class_label == 'explicit' or predicted_class_label == 'suggestive':
return json.dumps({"class": predicted_class_label})
return json.dumps({"class": "safe"})
except Exception as e:
print(f"Error in processing images: {str(e)}")
return json.dumps({"class": "safe"})
@app.route("/predict", methods=["GET"])
def predict_image():
try:
src = request.args.get("src")
# Download image from the provided URL
response = requests.get(src)
response.raise_for_status()
predicted_class_label = predict(response)
# Return the predictions
return json.dumps({"class": predicted_class_label})
except requests.exceptions.RequestException as e:
return json.dumps({"error": f"Request error: {str(e)}"})
except Exception as e:
return json.dumps({"error": f"An unexpected error occurred: {str(e)}"})
if __name__ == "__main__":
app.run(debug=True)
|