Add files

app.py

@@ -19,48 +19,53 @@ model = AutoModelForImageClassification.from_pretrained('carbon225/vit-base-patc
 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))
+    try:
+        # 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")
+        # 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
+        # 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]
+        # 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
+        return predicted_class_label
+    except Exception as e:
+        print(f"Error in predicting image: {str(e)}")
+        return None
 
 @app.route("/", methods=["GET"])
 def default():
     return json.dumps({"Server": "Working"})
 
-@app.route("/extractimages",methods=["GET"])
+@app.route("/extractimages", methods=["GET"])
 def extract_images():
     try:
-        src=request.args.get("src")
+        src = request.args.get("src")
         response = requests.get(src)
-        soup = BeautifulSoup(response.content,'html.parser')
+        soup = BeautifulSoup(response.content, 'html.parser')
 
         img_tags = soup.select('div img')
         for img_tag in img_tags:
             img_url = urljoin(src, img_tag['src'])
             response = requests.get(img_url)
-            response.raise_for_status()  
+            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})
+            if predicted_class_label == 'explicit' or predicted_class_label == 'suggestive':
+                return json.dumps({"class": predicted_class_label})
 
-        return json.dumps({"class":"safe"})
+        return json.dumps({"class": "safe"})
     except Exception as e:
-        return e
+        print(f"Error in processing images: {str(e)}")
+        return json.dumps({"class": "safe"})
 
 @app.route("/predict", methods=["GET"])
 def predict_image():
@@ -69,7 +74,7 @@ def predict_image():
 
         # Download image from the provided URL
         response = requests.get(src)
-        response.raise_for_status()  
+        response.raise_for_status()
 
         predicted_class_label = predict(response)