Spaces:
Sleeping
Sleeping
File size: 2,455 Bytes
7efe748 43d3f36 7efe748 c095153 7efe748 43d3f36 7efe748 c095153 095d14f c095153 095d14f c095153 7efe748 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 |
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
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 = 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() # Check for HTTP errors
# 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)
|