Skip to content

Commit

Permalink
feat: wrap decoding result for analysis
Browse files Browse the repository at this point in the history
  • Loading branch information
jandom committed May 25, 2024
1 parent e18267e commit 1c227f5
Showing 1 changed file with 65 additions and 6 deletions.
71 changes: 65 additions & 6 deletions qreader.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,10 @@
from warnings import warn
import numpy as np
from pyzbar.pyzbar import decode as decodeQR, ZBarSymbol, Decoded
from dataclasses import dataclass
import cv2
import os
import typing

from qrdet import QRDetector, crop_qr, PADDED_QUAD_XY, BBOX_XYXY, CONFIDENCE, CXCY, WH, POLYGON_XY, QUAD_XY

Expand All @@ -23,6 +25,34 @@
# In windows shift-jis is the default encoding will use, while in linux is big5
DEFAULT_REENCODINGS = ('shift-jis', 'big5') if os.name == 'nt' else ('big5', 'shift-jis')

@dataclass(frozen=True)
class DecodeQRResult():
scale_factor: float
corrections: typing.Literal["cropped_bbox", "corrected_perspective"]
flavor: typing.Literal["original", "inverted", "grayscale"]
blur_kernel_sizes: tuple[tuple[int, int]] | None
image: np.ndarray
result: Decoded

def wrap(
scale_factor: float,
corrections: str,
flavor: str,
blur_kernel_sizes: tuple[tuple[int, int]],
image: np.ndarray,
results: typing.List[Decoded],
) -> list[DecodeQRResult]:

return [DecodeQRResult(
scale_factor=scale_factor,
corrections=corrections,
flavor=flavor,
blur_kernel_sizes=blur_kernel_sizes,
image=image,
result=result,
) for result in results]


class QReader:
def __init__(self, model_size: str = 's', min_confidence: float = 0.5,
reencode_to: str | tuple[str] | list[str] | None = DEFAULT_REENCODINGS):
Expand Down Expand Up @@ -191,7 +221,7 @@ def _decode_qr_zbar(self, image: np.ndarray,
padded_quad_xy=updated_detection[PADDED_QUAD_XY])

for scale_factor in (1, 0.5, 2, 0.25, 3, 4):
for image in (cropped_bbox, corrected_perspective):
for label, image in {"cropped_bbox": cropped_bbox, "corrected_perspective": corrected_perspective}.items():
# If rescaled_image will be larger than 1024px, skip it
# TODO: Decide a minimum size for the QRs based on the resize benchmark
if not all(25 < axis < 1024 for axis in image.shape[:2]) and scale_factor != 1:
Expand All @@ -201,11 +231,26 @@ def _decode_qr_zbar(self, image: np.ndarray,
interpolation=cv2.INTER_CUBIC)
decodedQR = decodeQR(image=rescaled_image, symbols=[ZBarSymbol.QRCODE])
if len(decodedQR) > 0:
return decodedQR
return wrap(
scale_factor=scale_factor,
corrections=label,
flavor="original",
blur_kernel_sizes=None,
image=rescaled_image,
results=decodedQR
)
# For QRs with black background and white foreground, try to invert the image
decodedQR = decodeQR(image=255 - rescaled_image, symbols=[ZBarSymbol.QRCODE])
inverted_image = image=255 - rescaled_image
decodedQR = decodeQR(inverted_image, symbols=[ZBarSymbol.QRCODE])
if len(decodedQR) > 0:
return decodedQR
return wrap(
scale_factor=scale_factor,
corrections=label,
flavor="inverted",
blur_kernel_sizes=None,
image=inverted_image,
results=decodedQR,
)

# If it not works, try to parse to grayscale (if it is not already)
if len(rescaled_image.shape) == 3:
Expand All @@ -216,7 +261,14 @@ def _decode_qr_zbar(self, image: np.ndarray,
gray = rescaled_image
decodedQR = self.__threshold_and_blur_decodings(image=gray, blur_kernel_sizes=((5, 5), (7, 7)))
if len(decodedQR) > 0:
return decodedQR
return wrap(
scale_factor=scale_factor,
corrections=label,
flavor="grayscale",
blur_kernel_sizes=((5, 5), (7, 7)),
image=gray,
results=decodedQR
)

if len(rescaled_image.shape) == 3:
# If it not works, try to sharpen the image
Expand All @@ -226,7 +278,14 @@ def _decode_qr_zbar(self, image: np.ndarray,
sharpened_gray = cv2.filter2D(src=rescaled_image, ddepth=-1, kernel=_SHARPEN_KERNEL)
decodedQR = self.__threshold_and_blur_decodings(image=sharpened_gray, blur_kernel_sizes=((3, 3),))
if len(decodedQR) > 0:
return decodedQR
return wrap(
scale_factor=scale_factor,
corrections=label,
flavor="grayscale",
blur_kernel_sizes=((3, 3),),
image=sharpened_gray,
results=decodedQR
)

return []

Expand Down

0 comments on commit 1c227f5

Please sign in to comment.