Add files

app.py

@@ -18,6 +18,25 @@ cors = CORS(app)
 model = AutoModelForImageClassification.from_pretrained('carbon225/vit-base-patch16-224-hentai')
 feature_extractor = AutoFeatureExtractor.from_pretrained('carbon225/vit-base-patch16-224-hentai')
 
+def predict(response):
+    # 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 predicted_class_label
+
 @app.route("/", methods=["GET"])
 def default():
     return json.dumps({"Server": "Working"})
@@ -28,18 +47,23 @@ def extract_images():
         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})
+            response = requests.get(img_url)
+            response.raise_for_status()  
+            predicted_class_label = predict(response)
+
+            if predicted_class_label=='explicit' or predicted_class_label=='suggestive':
+                return json.dumps({"class":predicted_class_label})
+
+        return json.dumps({"class":"safe"})
     except Exception as e:
         return e
 
 @app.route("/predict", methods=["GET"])
-def predict():
+def predict_image():
     try:
         src = request.args.get("src")
 
@@ -47,21 +71,7 @@ def predict():
         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]
+        predicted_class_label = predict(response)
 
         # Return the predictions
         return json.dumps({"class": predicted_class_label})