Spaces:
Runtime error
Runtime error
Add files
Browse files
app.py
CHANGED
@@ -19,48 +19,53 @@ model = AutoModelForImageClassification.from_pretrained('carbon225/vit-base-patc
|
|
19 |
feature_extractor = AutoFeatureExtractor.from_pretrained('carbon225/vit-base-patch16-224-hentai')
|
20 |
|
21 |
def predict(response):
|
22 |
-
|
23 |
-
|
24 |
-
|
|
|
25 |
|
26 |
-
|
27 |
-
|
28 |
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
|
38 |
-
|
|
|
|
|
|
|
39 |
|
40 |
@app.route("/", methods=["GET"])
|
41 |
def default():
|
42 |
return json.dumps({"Server": "Working"})
|
43 |
|
44 |
-
@app.route("/extractimages",methods=["GET"])
|
45 |
def extract_images():
|
46 |
try:
|
47 |
-
src=request.args.get("src")
|
48 |
response = requests.get(src)
|
49 |
-
soup = BeautifulSoup(response.content,'html.parser')
|
50 |
|
51 |
img_tags = soup.select('div img')
|
52 |
for img_tag in img_tags:
|
53 |
img_url = urljoin(src, img_tag['src'])
|
54 |
response = requests.get(img_url)
|
55 |
-
response.raise_for_status()
|
56 |
predicted_class_label = predict(response)
|
57 |
|
58 |
-
if predicted_class_label=='explicit' or predicted_class_label=='suggestive':
|
59 |
-
return json.dumps({"class":predicted_class_label})
|
60 |
|
61 |
-
return json.dumps({"class":"safe"})
|
62 |
except Exception as e:
|
63 |
-
|
|
|
64 |
|
65 |
@app.route("/predict", methods=["GET"])
|
66 |
def predict_image():
|
@@ -69,7 +74,7 @@ def predict_image():
|
|
69 |
|
70 |
# Download image from the provided URL
|
71 |
response = requests.get(src)
|
72 |
-
response.raise_for_status()
|
73 |
|
74 |
predicted_class_label = predict(response)
|
75 |
|
|
|
19 |
feature_extractor = AutoFeatureExtractor.from_pretrained('carbon225/vit-base-patch16-224-hentai')
|
20 |
|
21 |
def predict(response):
|
22 |
+
try:
|
23 |
+
# Open and preprocess the image
|
24 |
+
image = Image.open(BytesIO(response.content))
|
25 |
+
image = image.resize((128, 128))
|
26 |
|
27 |
+
# Extract features using the pre-trained feature extractor
|
28 |
+
encoding = feature_extractor(images=image.convert("RGB"), return_tensors="pt")
|
29 |
|
30 |
+
# Make a prediction using the pre-trained model
|
31 |
+
with torch.no_grad():
|
32 |
+
outputs = model(**encoding)
|
33 |
+
logits = outputs.logits
|
34 |
|
35 |
+
# Get the predicted class index and label
|
36 |
+
predicted_class_idx = logits.argmax(-1).item()
|
37 |
+
predicted_class_label = model.config.id2label[predicted_class_idx]
|
38 |
|
39 |
+
return predicted_class_label
|
40 |
+
except Exception as e:
|
41 |
+
print(f"Error in predicting image: {str(e)}")
|
42 |
+
return None
|
43 |
|
44 |
@app.route("/", methods=["GET"])
|
45 |
def default():
|
46 |
return json.dumps({"Server": "Working"})
|
47 |
|
48 |
+
@app.route("/extractimages", methods=["GET"])
|
49 |
def extract_images():
|
50 |
try:
|
51 |
+
src = request.args.get("src")
|
52 |
response = requests.get(src)
|
53 |
+
soup = BeautifulSoup(response.content, 'html.parser')
|
54 |
|
55 |
img_tags = soup.select('div img')
|
56 |
for img_tag in img_tags:
|
57 |
img_url = urljoin(src, img_tag['src'])
|
58 |
response = requests.get(img_url)
|
59 |
+
response.raise_for_status()
|
60 |
predicted_class_label = predict(response)
|
61 |
|
62 |
+
if predicted_class_label == 'explicit' or predicted_class_label == 'suggestive':
|
63 |
+
return json.dumps({"class": predicted_class_label})
|
64 |
|
65 |
+
return json.dumps({"class": "safe"})
|
66 |
except Exception as e:
|
67 |
+
print(f"Error in processing images: {str(e)}")
|
68 |
+
return json.dumps({"class": "safe"})
|
69 |
|
70 |
@app.route("/predict", methods=["GET"])
|
71 |
def predict_image():
|
|
|
74 |
|
75 |
# Download image from the provided URL
|
76 |
response = requests.get(src)
|
77 |
+
response.raise_for_status()
|
78 |
|
79 |
predicted_class_label = predict(response)
|
80 |
|