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

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +7 -2
handler.py CHANGED
@@ -43,7 +43,8 @@ class EndpointHandler:
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)
@@ -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
- return {"error": str(e)}
 
 
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)}]