torinriley
commited on
Commit
•
b22dbd1
1
Parent(s):
cc3a466
Update handler.py
Browse files- handler.py +31 -31
handler.py
CHANGED
@@ -38,39 +38,39 @@ class EndpointHandler:
|
|
38 |
return image_tensor
|
39 |
|
40 |
def __call__(self, data):
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
|
48 |
-
|
49 |
-
|
50 |
|
51 |
-
|
52 |
-
|
53 |
|
54 |
-
|
55 |
-
|
56 |
-
|
57 |
|
58 |
-
|
59 |
-
|
60 |
-
|
61 |
-
|
62 |
-
|
63 |
-
|
64 |
-
|
65 |
-
|
66 |
-
|
67 |
-
|
68 |
-
|
69 |
-
|
70 |
-
|
71 |
-
|
72 |
-
|
73 |
|
74 |
-
|
75 |
-
|
76 |
-
|
|
|
38 |
return image_tensor
|
39 |
|
40 |
def __call__(self, data):
|
41 |
+
"""
|
42 |
+
Process incoming raw binary image data.
|
43 |
+
"""
|
44 |
+
try:
|
45 |
+
if "body" not in data:
|
46 |
+
return {"error": "No image data provided in request."}
|
47 |
|
48 |
+
image_bytes = data["body"]
|
49 |
+
image_tensor = self.preprocess_frame(image_bytes)
|
50 |
|
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:
|
61 |
+
x1, y1, x2, y2 = map(int, box)
|
62 |
+
label_text = self.label_map.get(label, "unknown")
|
63 |
+
results.append({
|
64 |
+
"label": label_text,
|
65 |
+
"score": round(score, 2),
|
66 |
+
"box": {
|
67 |
+
"xmin": x1,
|
68 |
+
"ymin": y1,
|
69 |
+
"xmax": x2,
|
70 |
+
"ymax": y2
|
71 |
+
}
|
72 |
+
})
|
73 |
|
74 |
+
return results
|
75 |
+
except Exception as e:
|
76 |
+
return {"error": str(e)}
|