File size: 1,614 Bytes
77753ce de5f953 2066911 77753ce addde62 de5f953 77753ce fc514c5 de5f953 28f318a 2066911 de5f953 28f318a de5f953 28f318a 36b4197 853070a cf572c0 db01789 cf572c0 36b4197 77753ce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
from typing import Union
from fastapi import FastAPI,File
from PIL import Image
from io import BytesIO
from transformers import DetrImageProcessor, DetrForObjectDetection
import torch
import requests
app = FastAPI(title="Object Detection",
docs_url="/",
description="Object detection in Image")
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
@app.post('/image')
def read_image(image_file: bytes = File(...)):
image = Image.open(BytesIO(image_file))
# url = "http://images.cocodataset.org/val2017/000000039769.jpg"
# image = Image.open(requests.get(url, stream=True).raw)
inputs = processor(images=image, return_tensors="pt")
print("image loaded")
outputs = model(**inputs)
target_sizes = torch.tensor([image.size[::-1]])
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
print("results pushed")
response = {}
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
box = [round(i, 2) for i in box.tolist()]
# response['scores'] = model.config.id2label[label.item()]
# response['labels'] = score.item()
# response['boxes'] = box
print(model.config.id2label[label.item()])
print(score.item())
print(box)
# print(
# f"Detected {model.config.id2label[label.item()]} with confidence "
# f"{round(score.item(), 3)} at location {box}"
# )
return response
|