HawkeyeHS commited on
Commit
918f252
1 Parent(s): 1dd2c0f
Files changed (3) hide show
  1. Dockerfile +15 -0
  2. app.py +75 -0
  3. requirements.txt +58 -0
Dockerfile ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.9
2
+
3
+ WORKDIR /code
4
+
5
+ # Create cache directory and set permissions
6
+ RUN mkdir -p /.cache/huggingface/hub && \
7
+ chmod -R 777 /.cache
8
+
9
+ COPY ./requirements.txt /code/requirements.txt
10
+
11
+ RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
12
+
13
+ COPY . .
14
+
15
+ CMD ["gunicorn","-b", "0.0.0.0:7860", "app:app"]
app.py ADDED
@@ -0,0 +1,75 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from transformers import AutoModelForImageClassification, AutoFeatureExtractor
3
+ import torch
4
+ from flask_cors import CORS
5
+ from flask import Flask, request, json
6
+ from PIL import Image
7
+ import requests
8
+ from io import BytesIO
9
+ from bs4 import BeautifulSoup
10
+ from urllib.parse import urljoin
11
+
12
+ os.environ["CUDA_VISIBLE_DEVICES"] = ""
13
+
14
+ app = Flask(__name__)
15
+ cors = CORS(app)
16
+
17
+ # Define the model and feature extractor globally
18
+ model = AutoModelForImageClassification.from_pretrained('carbon225/vit-base-patch16-224-hentai')
19
+ feature_extractor = AutoFeatureExtractor.from_pretrained('carbon225/vit-base-patch16-224-hentai')
20
+
21
+ @app.route("/", methods=["GET"])
22
+ def default():
23
+ return json.dumps({"Server": "Working"})
24
+
25
+ @app.route("/extractimages",methods=["GET"])
26
+ def extract_images():
27
+ try:
28
+ src=request.args.get("src")
29
+ response = requests.get(src)
30
+ soup = BeautifulSoup(response.content,'html.parser')
31
+ img_urls=[]
32
+
33
+ img_tags = soup.select('div img')
34
+ for img_tag in img_tags:
35
+ img_url = urljoin(src, img_tag['src'])
36
+ img_urls.append(img_url)
37
+ return json.dumps({"images":img_urls})
38
+ except Exception as e:
39
+ return e
40
+
41
+ @app.route("/predict", methods=["GET"])
42
+ def predict():
43
+ try:
44
+ src = request.args.get("src")
45
+
46
+ # Download image from the provided URL
47
+ response = requests.get(src)
48
+ response.raise_for_status()
49
+
50
+ # Open and preprocess the image
51
+ image = Image.open(BytesIO(response.content))
52
+ image = image.resize((128, 128))
53
+
54
+ # Extract features using the pre-trained feature extractor
55
+ encoding = feature_extractor(images=image.convert("RGB"), return_tensors="pt")
56
+
57
+ # Make a prediction using the pre-trained model
58
+ with torch.no_grad():
59
+ outputs = model(**encoding)
60
+ logits = outputs.logits
61
+
62
+ # Get the predicted class index and label
63
+ predicted_class_idx = logits.argmax(-1).item()
64
+ predicted_class_label = model.config.id2label[predicted_class_idx]
65
+
66
+ # Return the predictions
67
+ return json.dumps({"class": predicted_class_label})
68
+
69
+ except requests.exceptions.RequestException as e:
70
+ return json.dumps({"error": f"Request error: {str(e)}"})
71
+ except Exception as e:
72
+ return json.dumps({"error": f"An unexpected error occurred: {str(e)}"})
73
+
74
+ if __name__ == "__main__":
75
+ app.run(debug=True)
requirements.txt ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ blinker==1.7.0
2
+ bs4==0.0.1
3
+ certifi==2023.11.17
4
+ charset-normalizer==3.3.2
5
+ click==8.1.7
6
+ filelock==3.13.1
7
+ Flask==3.0.0
8
+ Flask-Cors==4.0.0
9
+ fsspec==2023.12.2
10
+ huggingface-hub==0.20.1
11
+ idna==3.6
12
+ importlib-metadata==7.0.0
13
+ itsdangerous==2.1.2
14
+ Jinja2==3.1.2
15
+ markdownify==0.11.6
16
+ MarkupSafe==2.1.3
17
+ mpmath==1.3.0
18
+ networkx==3.2.1
19
+ numpy==1.26.2
20
+ nvidia-cublas-cu12==12.1.3.1
21
+ nvidia-cuda-cupti-cu12==12.1.105
22
+ nvidia-cuda-nvrtc-cu12==12.1.105
23
+ nvidia-cuda-runtime-cu12==12.1.105
24
+ nvidia-cudnn-cu12==8.9.2.26
25
+ nvidia-cufft-cu12==11.0.2.54
26
+ nvidia-curand-cu12==10.3.2.106
27
+ nvidia-cusolver-cu12==11.4.5.107
28
+ nvidia-cusparse-cu12==12.1.0.106
29
+ nvidia-nccl-cu12==2.18.1
30
+ nvidia-nvjitlink-cu12==12.3.101
31
+ nvidia-nvtx-cu12==12.1.105
32
+ outcome==1.2.0
33
+ packaging==23.2
34
+ perplexity-api @ git+https://gitlab.com/robowaifudev/perplexity-api@0baebed39beeedc13e6281af5505f4c4221fdb5f
35
+ Pillow==10.1.0
36
+ PySocks==1.7.1
37
+ PyYAML==6.0.1
38
+ regex==2023.10.3
39
+ requests==2.31.0
40
+ safetensors==0.4.1
41
+ selenium==4.12.0
42
+ sortedcontainers==2.4.0
43
+ sympy==1.12
44
+ tokenizers==0.15.0
45
+ torch==2.1.2
46
+ tqdm==4.66.1
47
+ transformers==4.36.2
48
+ trio==0.22.2
49
+ trio-websocket==0.10.4
50
+ triton==2.1.0
51
+ typing_extensions==4.9.0
52
+ urllib3==2.1.0
53
+ Werkzeug==3.0.1
54
+ wsproto==1.2.0
55
+ zipp==3.17.0
56
+ gunicorn
57
+ torch
58
+ bs4