sarim commited on
Commit
de5f953
·
1 Parent(s): c65b181

image object detection

Browse files
Files changed (1) hide show
  1. main.py +22 -5
main.py CHANGED
@@ -1,14 +1,31 @@
1
  from typing import Union
2
 
3
- from fastapi import FastAPI
 
 
 
 
 
 
4
 
5
  app = FastAPI(title="ReceiptOCR",
6
- docs_url="/", description="Object detection in Image")
 
 
 
 
 
 
7
 
8
 
9
 
 
 
 
 
 
 
 
 
10
 
11
 
12
- @app.get("/items/{item_id}")
13
- def read_item(item_id: int, q: Union[str, None] = None):
14
- return {"item_id": item_id, "q": q}
 
1
  from typing import Union
2
 
3
+ from fastapi import FastAPI,File
4
+ from PIL import Image
5
+ from transformers import pipeline
6
+ from io import BytesIO
7
+ from transformers import DetrImageProcessor, DetrForObjectDetection
8
+ import torch
9
+
10
 
11
  app = FastAPI(title="ReceiptOCR",
12
+ docs_url="/",
13
+ title = 'Object Detection',
14
+ description="Object detection in Image")
15
+
16
+
17
+ processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
18
+ model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
19
 
20
 
21
 
22
+ @app.post('/image')
23
+ def read_image(image_file: bytes = File(...)):
24
+ image = Image.open(BytesIO(image_file))
25
+ inputs = processor(images=image, return_tensors="pt")
26
+ outputs = model(**inputs)
27
+ target_sizes = torch.tensor([image.size[::-1]])
28
+ results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
29
+ return results
30
 
31