Apply suggestions from @ThomasDh-C code review
Browse files- handler.py +6 -16
handler.py
CHANGED
|
@@ -24,11 +24,9 @@ easyocr.Reader(["en"])
|
|
| 24 |
class EndpointHandler:
|
| 25 |
def __init__(self, model_dir: str = "/repository") -> None:
|
| 26 |
self.device = (
|
| 27 |
-
torch.device("cuda")
|
| 28 |
-
if torch.
|
| 29 |
-
|
| 30 |
-
if torch.backends.mps.is_available()
|
| 31 |
-
else torch.device("cpu")
|
| 32 |
)
|
| 33 |
|
| 34 |
# bounding box detection model
|
|
@@ -54,7 +52,7 @@ class EndpointHandler:
|
|
| 54 |
# data should contain the following:
|
| 55 |
# "inputs": {
|
| 56 |
# "image": url/base64,
|
| 57 |
-
# (optional) "image_size":
|
| 58 |
# (optional) "bbox_threshold": float,
|
| 59 |
# (optional) "iou_threshold": float,
|
| 60 |
# }
|
|
@@ -63,25 +61,17 @@ class EndpointHandler:
|
|
| 63 |
# read image from either url or base64 encoding
|
| 64 |
image = load_image(data["image"])
|
| 65 |
|
| 66 |
-
# box_overlay_ratio = image.size[0] / 3200
|
| 67 |
-
# bbox_config = {
|
| 68 |
-
# "text_scale": 0.8 * box_overlay_ratio,
|
| 69 |
-
# "text_thickness": max(int(2 * box_overlay_ratio), 1),
|
| 70 |
-
# "text_padding": max(int(3 * box_overlay_ratio), 1),
|
| 71 |
-
# "thickness": max(int(3 * box_overlay_ratio), 1),
|
| 72 |
-
# }
|
| 73 |
-
|
| 74 |
ocr_texts, ocr_bboxes = self.check_ocr_bboxes(
|
| 75 |
image,
|
| 76 |
out_format="xyxy",
|
| 77 |
-
ocr_kwargs={"
|
| 78 |
)
|
| 79 |
annotated_image, filtered_bboxes_out = self.get_som_labeled_img(
|
| 80 |
image,
|
| 81 |
image_size=data.get("image_size", None),
|
| 82 |
ocr_texts=ocr_texts,
|
| 83 |
ocr_bboxes=ocr_bboxes,
|
| 84 |
-
bbox_threshold=data.get("bbox_threshold", 0.
|
| 85 |
iou_threshold=data.get("iou_threshold", None),
|
| 86 |
)
|
| 87 |
return {
|
|
|
|
| 24 |
class EndpointHandler:
|
| 25 |
def __init__(self, model_dir: str = "/repository") -> None:
|
| 26 |
self.device = (
|
| 27 |
+
torch.device("cuda") if torch.cuda.is_available()
|
| 28 |
+
else (torch.device("mps") if torch.backends.mps.is_available()
|
| 29 |
+
else torch.device("cpu"))
|
|
|
|
|
|
|
| 30 |
)
|
| 31 |
|
| 32 |
# bounding box detection model
|
|
|
|
| 52 |
# data should contain the following:
|
| 53 |
# "inputs": {
|
| 54 |
# "image": url/base64,
|
| 55 |
+
# (optional) "image_size": {"w": int, "h": int},
|
| 56 |
# (optional) "bbox_threshold": float,
|
| 57 |
# (optional) "iou_threshold": float,
|
| 58 |
# }
|
|
|
|
| 61 |
# read image from either url or base64 encoding
|
| 62 |
image = load_image(data["image"])
|
| 63 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 64 |
ocr_texts, ocr_bboxes = self.check_ocr_bboxes(
|
| 65 |
image,
|
| 66 |
out_format="xyxy",
|
| 67 |
+
ocr_kwargs={"text_threshold": 0.8},
|
| 68 |
)
|
| 69 |
annotated_image, filtered_bboxes_out = self.get_som_labeled_img(
|
| 70 |
image,
|
| 71 |
image_size=data.get("image_size", None),
|
| 72 |
ocr_texts=ocr_texts,
|
| 73 |
ocr_bboxes=ocr_bboxes,
|
| 74 |
+
bbox_threshold=data.get("bbox_threshold", 0.05),
|
| 75 |
iou_threshold=data.get("iou_threshold", None),
|
| 76 |
)
|
| 77 |
return {
|