torinriley
commited on
Commit
•
f132f1f
1
Parent(s):
b22dbd1
Update handler.py
Browse files- handler.py +7 -2
handler.py
CHANGED
@@ -43,7 +43,8 @@ class EndpointHandler:
|
|
43 |
"""
|
44 |
try:
|
45 |
if "body" not in data:
|
46 |
-
|
|
|
47 |
|
48 |
image_bytes = data["body"]
|
49 |
image_tensor = self.preprocess_frame(image_bytes)
|
@@ -51,10 +52,12 @@ class EndpointHandler:
|
|
51 |
with torch.no_grad():
|
52 |
predictions = self.model(image_tensor)
|
53 |
|
|
|
54 |
boxes = predictions[0]["boxes"].cpu().tolist()
|
55 |
labels = predictions[0]["labels"].cpu().tolist()
|
56 |
scores = predictions[0]["scores"].cpu().tolist()
|
57 |
|
|
|
58 |
results = []
|
59 |
for box, label, score in zip(boxes, labels, scores):
|
60 |
if score >= CONFIDENCE_THRESHOLD:
|
@@ -71,6 +74,8 @@ class EndpointHandler:
|
|
71 |
}
|
72 |
})
|
73 |
|
|
|
74 |
return results
|
75 |
except Exception as e:
|
76 |
-
|
|
|
|
43 |
"""
|
44 |
try:
|
45 |
if "body" not in data:
|
46 |
+
# Return error in the expected output format
|
47 |
+
return [{"error": "No image data provided in request."}]
|
48 |
|
49 |
image_bytes = data["body"]
|
50 |
image_tensor = self.preprocess_frame(image_bytes)
|
|
|
52 |
with torch.no_grad():
|
53 |
predictions = self.model(image_tensor)
|
54 |
|
55 |
+
# Extract predictions
|
56 |
boxes = predictions[0]["boxes"].cpu().tolist()
|
57 |
labels = predictions[0]["labels"].cpu().tolist()
|
58 |
scores = predictions[0]["scores"].cpu().tolist()
|
59 |
|
60 |
+
# Build the results array
|
61 |
results = []
|
62 |
for box, label, score in zip(boxes, labels, scores):
|
63 |
if score >= CONFIDENCE_THRESHOLD:
|
|
|
74 |
}
|
75 |
})
|
76 |
|
77 |
+
# Return results in the required schema
|
78 |
return results
|
79 |
except Exception as e:
|
80 |
+
# Return errors wrapped in a list to match the required schema
|
81 |
+
return [{"error": str(e)}]
|