torinriley commited on
Commit
b22dbd1
1 Parent(s): cc3a466

Update handler.py

Browse files
Files changed (1) hide show
  1. 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
- 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)}
 
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)}