Update notebook and inference script

Author: MalloryWittwer

inference.py

@@ -12,6 +12,7 @@
 from skimage.transform import rescale, resize
 from skimage.util import img_as_ubyte
 from skimage.exposure import rescale_intensity
+from skimage.segmentation import clear_border
 
 from preprocessing import keep_n_biggest_objects
 
@@ -20,12 +21,16 @@
     pipe = load(f)
 
 
-def extract_digit_from_digit_square(image: np.ndarray) -> np.ndarray:
+def extract_digit_from_digit_roi(image: np.ndarray) -> np.ndarray:
     image = img_as_ubyte(image)
 
     resized = rescale(image, scale=0.5)
 
-    binary = resized > threshold_otsu(resized)
+    try:
+        binary = resized > threshold_otsu(resized)
+    except ValueError:
+        return img_as_ubyte(np.zeros((50, 50)))
+    
     labelled = label(binary)
 
     light_square = keep_n_biggest_objects(labelled, n=1)
@@ -54,20 +59,26 @@ def extract_digit_from_digit_square(image: np.ndarray) -> np.ndarray:
     return digit_crop
 
 
-def detect_digit_squares(image: np.ndarray, n_squares=12) -> pd.DataFrame:
+def detect_digit_roi(image: np.ndarray, max_objects=12) -> pd.DataFrame:
+    # RGB => Grayscale
     gray = rgb2gray(image)
-    binary = gray < threshold_otsu(gray)
+    
+    # Otsu threshold
+    binary = gray > threshold_otsu(gray)
+    
+    # Labelling + keep biggest object
     labelled = label(binary)
-    grid_frame = keep_n_biggest_objects(labelled, n=1)
-    filled_grid_frame = ndi.binary_fill_holes(grid_frame)
-    digit_squares = np.logical_and(filled_grid_frame, grid_frame == 0)
-    digit_squares_labelled = label(digit_squares)
-    digit_squares_cleaned = keep_n_biggest_objects(digit_squares_labelled, n=n_squares)
+    
+    # Remove object touching the image border (incl. the white background, probably)
+    labelled = clear_border(labelled)
+    
+    # Keep the biggest `max_objects` objects only
+    digit_crops = keep_n_biggest_objects(labelled, n=max_objects)
 
     # Return a dataframe with intensity image and centroids
     df = pd.DataFrame(
         regionprops_table(
-            digit_squares_cleaned,
+            digit_crops,
             intensity_image=gray,
             properties=["intensity_image", "bbox"],
         )
@@ -77,40 +88,41 @@ def detect_digit_squares(image: np.ndarray, n_squares=12) -> pd.DataFrame:
 
 
 def classify_digit_square(digit_square_img: np.ndarray) -> str:
-    digit = extract_digit_from_digit_square(digit_square_img)
+    digit = extract_digit_from_digit_roi(digit_square_img)
 
     # Normalize intensity values
     digit_normed = rescale_intensity(digit, out_range=(0, 1))
 
+    # Flatten pixels
     pixel_features = np.reshape(digit_normed[None], (1, -1))
+    
+    # Predict
     pred_label = pipe.predict(pixel_features)[0]
 
     return pred_label
 
 
-def draw_boxes(img, df: pd.DataFrame):
+def draw_boxes(img, df: pd.DataFrame, font_size=1.5, thickness=2):
     boxes = df[['bbox-0', 'bbox-1', 'bbox-2', 'bbox-3']].values
     labels = df["pred"].values
 
     output = img.copy()
     for i, box in enumerate(boxes):
         y1, x1, y2, x2 = map(int, box)
 
-        cv2.rectangle(output, (x1, y1), (x2, y2), color=(0, 255, 0), thickness=2)
+        cv2.rectangle(output, (x1, y1), (x2, y2), color=(0, 255, 0), thickness=thickness)
 
         text = str(labels[i])
 
         (text_w, text_h), baseline = cv2.getTextSize(
             text=text, 
             fontFace=cv2.FONT_HERSHEY_SIMPLEX, 
-            fontScale=0.5, 
-            thickness=1,
+            fontScale=font_size, 
+            thickness=thickness,
         )
 
-        # Put label above box if there is room, otherwise inside/top edge
-        text_y = y1 - 8 if y1 - 8 > text_h else y1 + text_h + 8
+        text_y = y1 + text_h + 14
 
-        # Background rectangle for text
         cv2.rectangle(
             output,
             (x1, text_y - text_h - baseline),
@@ -124,9 +136,9 @@ def draw_boxes(img, df: pd.DataFrame):
             text=text,
             org=(x1, text_y),
             fontFace=cv2.FONT_HERSHEY_SIMPLEX,
-            fontScale=0.5,
+            fontScale=font_size,
             color=(0, 0, 0),
-            thickness=1,
+            thickness=thickness,
             lineType=cv2.LINE_AA,
         )
 
@@ -139,13 +151,15 @@ def main(camera_idx: int):
     if not cap.isOpened():
         raise RuntimeError("Could not open camera")
 
+    cv2.namedWindow("Camera", cv2.WINDOW_NORMAL)
+    
     while True:
         ret, frame = cap.read()
         if not ret:
             print("Failed to grab frame")
             break
 
-        df = detect_digit_squares(frame)
+        df = detect_digit_roi(frame)
 
         df["pred"] = df["intensity_image"].apply(classify_digit_square)