IHappyPlant commited on
Commit
c92f607
1 Parent(s): 03ef2e6
Files changed (1) hide show
  1. app.py +6 -4
app.py CHANGED
@@ -1,4 +1,5 @@
1
  from io import BytesIO
 
2
 
3
  from fastapi import FastAPI, Response, status, UploadFile
4
  from torchvision.io import read_image
@@ -24,20 +25,21 @@ async def healthcheck():
24
  return Response(status_code=status.HTTP_200_OK)
25
 
26
 
27
- @app.post('/detectObjectsFromURL')
28
  async def infer(image: UploadFile):
29
- img = read_image(image.filename)
 
 
 
30
 
31
  batch = [preprocess(img)]
32
 
33
- # Step 4: Use the model and visualize the prediction
34
  prediction = model(batch)[0]
35
  labels = [weights.meta["categories"][i] for i in prediction["labels"]]
36
  box = draw_bounding_boxes(img, boxes=prediction["boxes"],
37
  labels=labels,
38
  colors="red",
39
  width=4, font_size=30)
40
- im: Image.Image
41
  im = to_pil_image(box.detach())
42
  with BytesIO() as bio:
43
  im.save(bio, format='PNG')
 
1
  from io import BytesIO
2
+ from tempfile import NamedTemporaryFile
3
 
4
  from fastapi import FastAPI, Response, status, UploadFile
5
  from torchvision.io import read_image
 
25
  return Response(status_code=status.HTTP_200_OK)
26
 
27
 
28
+ @app.post('/detectObjects')
29
  async def infer(image: UploadFile):
30
+ with NamedTemporaryFile() as f:
31
+ f.write(image.file.read())
32
+ f.seek(0)
33
+ img = read_image(f.name)
34
 
35
  batch = [preprocess(img)]
36
 
 
37
  prediction = model(batch)[0]
38
  labels = [weights.meta["categories"][i] for i in prediction["labels"]]
39
  box = draw_bounding_boxes(img, boxes=prediction["boxes"],
40
  labels=labels,
41
  colors="red",
42
  width=4, font_size=30)
 
43
  im = to_pil_image(box.detach())
44
  with BytesIO() as bio:
45
  im.save(bio, format='PNG')