HawkeyeHS commited on
Commit
c095153
·
1 Parent(s): 43d3f36

Add extract image endpoint

Browse files
Files changed (1) hide show
  1. app.py +18 -4
app.py CHANGED
@@ -8,6 +8,7 @@ import numpy as np
8
  from PIL import Image
9
  import requests
10
  from io import BytesIO
 
11
 
12
  os.environ["CUDA_VISIBLE_DEVICES"] = ""
13
 
@@ -20,13 +21,28 @@ feature_extractor = AutoFeatureExtractor.from_pretrained('carbon225/vit-base-pat
20
 
21
  @app.route("/", methods=["GET"])
22
  def default():
23
- return json.dumps({"Hello I am Chitti": "Speed 1 Terra Hertz, Memory 1 Zeta Byte"})
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
 
25
  @app.route("/predict", methods=["GET"])
26
  def predict():
27
  try:
28
  src = request.args.get("src")
29
- print(f"{src=}")
30
 
31
  # Download image from the provided URL
32
  response = requests.get(src)
@@ -48,8 +64,6 @@ def predict():
48
  predicted_class_idx = logits.argmax(-1).item()
49
  predicted_class_label = model.config.id2label[predicted_class_idx]
50
 
51
- print(predicted_class_label)
52
-
53
  # Return the predictions
54
  return json.dumps({"class": predicted_class_label})
55
 
 
8
  from PIL import Image
9
  import requests
10
  from io import BytesIO
11
+ from bs4 import BeautifulSoup
12
 
13
  os.environ["CUDA_VISIBLE_DEVICES"] = ""
14
 
 
21
 
22
  @app.route("/", methods=["GET"])
23
  def default():
24
+ return json.dumps({"Server": "Working"})
25
+
26
+ @app.route("/extractimages",methods=["GET"])
27
+ def extract_images():
28
+ try:
29
+ src=request.args.get("src")
30
+ response = requests.get(src)
31
+ soup = BeautifulSoup(response.text, 'html.parser')
32
+ img_urls=[]
33
+
34
+ img_tags = soup.find_all('img')
35
+ for img_tag in img_tags:
36
+ img_url = img_tag['src']
37
+ img_urls.append(img_url)
38
+ return json.dumps({"images":img_urls})
39
+ except Exception as e:
40
+ return e
41
 
42
  @app.route("/predict", methods=["GET"])
43
  def predict():
44
  try:
45
  src = request.args.get("src")
 
46
 
47
  # Download image from the provided URL
48
  response = requests.get(src)
 
64
  predicted_class_idx = logits.argmax(-1).item()
65
  predicted_class_label = model.config.id2label[predicted_class_idx]
66
 
 
 
67
  # Return the predictions
68
  return json.dumps({"class": predicted_class_label})
69