|
|
import os |
|
|
import torch |
|
|
from torchvision import transforms |
|
|
from PIL import Image |
|
|
import io |
|
|
|
|
|
from model import get_model |
|
|
|
|
|
class EndpointHandler: |
|
|
def __init__(self, path: str = ""): |
|
|
""" |
|
|
Initialize the handler. Load the Faster R-CNN model. |
|
|
""" |
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
self.model_weights_path = os.path.join(path, "model.pt") |
|
|
|
|
|
|
|
|
self.model = get_model(num_classes=4) |
|
|
print(f"Loading weights from: {self.model_weights_path}") |
|
|
checkpoint = torch.load(self.model_weights_path, map_location=self.device) |
|
|
self.model.load_state_dict(checkpoint["model_state_dict"]) |
|
|
self.model.to(self.device) |
|
|
self.model.eval() |
|
|
|
|
|
|
|
|
self.transform = transforms.Compose([ |
|
|
transforms.Resize((640, 640)), |
|
|
transforms.ToTensor(), |
|
|
]) |
|
|
|
|
|
def __call__(self, data): |
|
|
""" |
|
|
Process the incoming request and return object detection predictions. |
|
|
""" |
|
|
try: |
|
|
if "image" not in data: |
|
|
return [{"error": "No 'image' provided in request."}] |
|
|
|
|
|
image_bytes = data["image"].encode("latin1") |
|
|
image = Image.open(io.BytesIO(image_bytes)).convert("RGB") |
|
|
|
|
|
input_tensor = self.transform(image).unsqueeze(0).to(self.device) |
|
|
|
|
|
with torch.no_grad(): |
|
|
outputs = self.model(input_tensor) |
|
|
|
|
|
boxes = outputs[0]["boxes"].cpu().tolist() |
|
|
labels = outputs[0]["labels"].cpu().tolist() |
|
|
scores = outputs[0]["scores"].cpu().tolist() |
|
|
|
|
|
threshold = 0.5 |
|
|
predictions = [ |
|
|
{"box": box, "label": label, "score": score} |
|
|
for box, label, score in zip(boxes, labels, scores) |
|
|
if score > threshold |
|
|
] |
|
|
|
|
|
return [{"predictions": predictions}] |
|
|
except Exception as e: |
|
|
return [{"error": str(e)}] |
