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()  # 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)