Add extract image endpoint

app.py

@@ -8,6 +8,7 @@ import numpy as np
 from PIL import Image
 import requests
 from io import BytesIO
+from bs4 import BeautifulSoup
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
@@ -20,13 +21,28 @@ feature_extractor = AutoFeatureExtractor.from_pretrained('carbon225/vit-base-pat
 
 @app.route("/", methods=["GET"])
 def default():
-    return json.dumps({"Hello I am Chitti": "Speed 1 Terra Hertz, Memory 1 Zeta Byte"})
+    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.text, 'html.parser')
+        img_urls=[]
+
+        img_tags = soup.find_all('img')
+        for img_tag in img_tags:
+            img_url = 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")
-        print(f"{src=}")
 
         # Download image from the provided URL
         response = requests.get(src)
@@ -48,8 +64,6 @@ def predict():
         predicted_class_idx = logits.argmax(-1).item()
         predicted_class_label = model.config.id2label[predicted_class_idx]
 
-        print(predicted_class_label)
-
         # Return the predictions
         return json.dumps({"class": predicted_class_label})