Spaces:
Sleeping
Sleeping
File size: 2,479 Bytes
7efe748 43d3f36 7efe748 c095153 9636ab8 7efe748 43d3f36 7efe748 c095153 095d14f c095153 095d14f c095153 9636ab8 c095153 7efe748 43d3f36 6c56640 43d3f36 7efe748 43d3f36 7efe748 43d3f36 7efe748 43d3f36 a8dea39 43d3f36 a8dea39 43d3f36 |
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 |
import os
import warnings
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
import torch
from flask_cors import CORS
from flask import Flask, request, json, Response
import numpy as np
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')
@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_urls=[]
img_tags = soup.select('div img')
for img_tag in img_tags:
img_url = urljoin(src, img_tag['src'])
img_urls.append(img_url)
return json.dumps({"images":img_urls})
except Exception as e:
return e
@app.route("/predict", methods=["GET"])
def predict():
try:
src = request.args.get("src")
# Download image from the provided URL
response = requests.get(src)
response.raise_for_status()
# 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 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)
|