Initial commit.

.flake8

@@ -0,0 +1,5 @@
+[flake8]
+ignore = I001,I002,I004,I005,I101,I201,C101,C403,C901,F401,F403,S001,D100,D101,D102,D103,D104,D105,D106,D107,D200,D210,D205,D400,T001,W504,D202,E203,W503,B006,T002,T100,P103,C408,F841
+max-line-length = 120
+exclude = outputs/*
+max-complexity = 10

.gitattributes

@@ -1,2 +1,3 @@
 # Auto detect text files and perform LF normalization
 * text=auto
+*.pth filter=lfs diff=lfs merge=lfs -text

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.pre-commit-config.yaml

@@ -0,0 +1,35 @@
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.3.0
+    hooks:
+    -   id: check-yaml
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+-   repo: https://github.com/psf/black
+    rev: '22.12.0'
+    hooks:
+    -   id: black
+        args: [--config=pyproject.toml]
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: 586b4f0
+    hooks:
+    - id: mypy
+      args: [--ignore-missing-imports, --warn-no-return, --warn-redundant-casts, --disallow-incomplete-defs, --no-namespace-packages ]
+-   repo: https://gitlab.com/pycqa/flake8
+    rev: '5.0.4'
+    hooks:
+    -   id: flake8
+        additional_dependencies: [
+                'flake8-bugbear==22.8.23',
+                'flake8-coding==1.3.2',
+                'flake8-comprehensions==3.10.0',
+                'flake8-debugger==4.1.2',
+                'flake8-deprecated==1.3',
+                'flake8-docstrings==1.6.0',
+                'flake8-isort==4.2.0',
+                'flake8-pep3101==1.3.0',
+                'flake8-polyfill==1.0.2',
+                'flake8-print==5.0.0',
+                'flake8-quotes==3.3.1',
+                'flake8-string-format==0.3.0',
+        ]

README.md

@@ -1,2 +1,17 @@
 # digit-draw-detect
  An app for handwritten digit detection
+
+steps:
+* use git lfs for the model +
+* write better code +
+* convert model to onnx or some other format?
+* deploy bare working app, without nice things
+* make better design
+* think about descriptions on the site
+
+On using git lfs:
+```shell
+git lfs install
+git lfs track "*.psd"
+git add .gitattributes
+```

config.toml

@@ -0,0 +1,8 @@
+# Everything in this section will be available as an environment variable
+db_username = "Jane"
+db_password = "12345qwerty"
+
+AWS_ACCESS_KEY_ID = 'AKIAI4JDKBYRCHGT77VQ'
+AWS_SECRET_ACCESS_KEY = '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'
+S3_BUCKET = 'digitdrawdetect'
+S3_BUCKET_NAME = 'digitdrawdetect'

model_files/best_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:571e937122d5ccafe496d1cc71cea5c0661d385b5a7db4ec977ac8ae5da40680
+size 246698572

mypy.ini

@@ -0,0 +1,5 @@
+import numpy
+
+[mypy]
+python_version = 3.10
+plugins = numpy.typing.mypy_plugin

pyproject.toml

@@ -0,0 +1,20 @@
+[tool.black]
+line-length = 119
+skip-string-normalization = true
+target-version = ['py39', 'py310']
+include = '\.pyi?$'
+exclude = '''
+/(
+    \.eggs
+  | \.git
+  | \.hg
+  | \.mypy_cache
+  | \.tox
+  | \.venv
+  | _build
+  | buck-out
+  | build
+  | dist
+  | outputs
+)/
+'''

requirements.txt

@@ -0,0 +1,15 @@
+albumentations==1.3.0
+matplotlib==3.6.1
+numpy==1.23.4
+omegaconf==2.2.1
+opencv_python==4.6.0.66
+pandas==1.5.1
+Pillow==9.2.0
+rich==12.6.0
+streamlit==1.16.0
+streamlit_drawable_canvas==0.9.2
+tomli==2.0.1
+torch==1.12.1
+torchvision
+# need to define some pytorch
+https://download.pytorch.org/whl/cpu/torch-1.13.1%2Bcpu-cp310-cp310-linux_x86_64.whl

src/ml_utils.py

@@ -0,0 +1,207 @@
+from typing import List
+import logging
+
+import albumentations as A
+import streamlit as st
+import torch
+from albumentations import pytorch
+
+from src.model_architecture import Net
+
+anchors = torch.tensor(
+    [
+        [[0.2800, 0.2200], [0.3800, 0.4800], [0.9000, 0.7800]],
+        [[0.0700, 0.1500], [0.1500, 0.1100], [0.1400, 0.2900]],
+        [[0.0200, 0.0300], [0.0400, 0.0700], [0.0800, 0.0600]],
+    ]
+)
+
+transforms = A.Compose(
+    [
+        A.Resize(always_apply=False, p=1, height=192, width=192, interpolation=1),
+        A.Normalize(),
+        pytorch.transforms.ToTensorV2(),
+    ]
+)
+
+
+def cells_to_bboxes(predictions: torch.tensor, anchors: torch.tensor, s: int, is_preds: bool = True) -> torch.tensor:
+    """
+    Scale the predictions coming from the model_files to
+    be relative to the entire image such that they for example later
+    can be plotted or.
+    Args:
+        predictions: tensor of size (N, 3, S, S, num_classes+5)
+        anchors: the anchors used for the predictions
+        s: the number of cells the image is divided in on the width (and height)
+        is_preds: whether the input is predictions or the true bounding boxes
+    Returns:
+        converted_bboxes: the converted boxes of sizes (N, num_anchors, S, S, 1+5) with class index,
+                      object score, bounding box coordinates
+    """
+    batch_size = predictions.shape[0]
+    num_anchors = len(anchors)
+    box_predictions = predictions[..., 1:5]
+    if is_preds:
+        anchors = anchors.reshape(1, len(anchors), 1, 1, 2)
+        box_predictions[..., 0:2] = torch.sigmoid(box_predictions[..., 0:2])
+        box_predictions[..., 2:] = torch.exp(box_predictions[..., 2:]) * anchors
+        scores = torch.sigmoid(predictions[..., 0:1])
+        best_class = torch.argmax(predictions[..., 5:], dim=-1).unsqueeze(-1)
+    else:
+        scores = predictions[..., 0:1]
+        best_class = predictions[..., 5:6]
+
+    cell_indices = torch.arange(s).repeat(predictions.shape[0], 3, s, 1).unsqueeze(-1).to(predictions.device)
+    x = 1 / s * (box_predictions[..., 0:1] + cell_indices)
+    y = 1 / s * (box_predictions[..., 1:2] + cell_indices.permute(0, 1, 3, 2, 4))
+    w_h = 1 / s * box_predictions[..., 2:4]
+    converted_bboxes = torch.cat((best_class, scores, x, y, w_h), dim=-1).reshape(batch_size, num_anchors * s * s, 6)
+    return converted_bboxes.tolist()
+
+
+def non_max_suppression(
+    bboxes: List[List], iou_threshold: float, threshold: float, box_format: str = 'corners'
+) -> List[List]:
+    """
+    Apply nms to the bboxes.
+
+    Video explanation of this function:
+    https://youtu.be/YDkjWEN8jNA
+    Does Non Max Suppression given bboxes
+    Args:
+        bboxes (list): list of lists containing all bboxes with each bboxes
+        specified as [class_pred, prob_score, x1, y1, x2, y2]
+        iou_threshold (float): threshold where predicted bboxes is correct
+        threshold (float): threshold to remove predicted bboxes (independent of IoU)
+        box_format (str): 'midpoint' or 'corners' used to specify bboxes
+    Returns:
+        list: bboxes after performing NMS given a specific IoU threshold
+    """
+
+    assert type(bboxes) == list
+
+    bboxes = [box for box in bboxes if box[1] > threshold]
+    bboxes = sorted(bboxes, key=lambda x: x[1], reverse=True)
+    bboxes_after_nms = []
+
+    while bboxes:
+        chosen_box = bboxes.pop(0)
+
+        bboxes = [
+            box
+            for box in bboxes
+            if box[0] != chosen_box[0]
+            or intersection_over_union(
+                torch.tensor(chosen_box[2:]),
+                torch.tensor(box[2:]),
+                box_format=box_format,
+            )
+            < iou_threshold
+        ]
+
+        bboxes_after_nms.append(chosen_box)
+
+    return bboxes_after_nms
+
+
+def intersection_over_union(
+    boxes_preds: torch.tensor, boxes_labels: torch.tensor, box_format: str = 'midpoint'
+) -> torch.tensor:
+    """
+    Calculate iou.
+
+    Video explanation of this function:
+    https://youtu.be/XXYG5ZWtjj0
+    This function calculates intersection over union (iou) given pred boxes
+    and target boxes.
+    Args:
+        boxes_preds (tensor): Predictions of Bounding Boxes (BATCH_SIZE, 4)
+        boxes_labels (tensor): Correct labels of Bounding Boxes (BATCH_SIZE, 4)
+        box_format (str): midpoint/corners, if boxes (x,y,w,h) or (x1,y1,x2,y2)
+    Returns:
+        tensor: Intersection over union for all examples
+    """
+
+    if box_format == 'midpoint':
+        box1_x1 = boxes_preds[..., 0:1] - boxes_preds[..., 2:3] / 2
+        box1_y1 = boxes_preds[..., 1:2] - boxes_preds[..., 3:4] / 2
+        box1_x2 = boxes_preds[..., 0:1] + boxes_preds[..., 2:3] / 2
+        box1_y2 = boxes_preds[..., 1:2] + boxes_preds[..., 3:4] / 2
+        box2_x1 = boxes_labels[..., 0:1] - boxes_labels[..., 2:3] / 2
+        box2_y1 = boxes_labels[..., 1:2] - boxes_labels[..., 3:4] / 2
+        box2_x2 = boxes_labels[..., 0:1] + boxes_labels[..., 2:3] / 2
+        box2_y2 = boxes_labels[..., 1:2] + boxes_labels[..., 3:4] / 2
+
+    if box_format == 'corners':
+        box1_x1 = boxes_preds[..., 0:1]
+        box1_y1 = boxes_preds[..., 1:2]
+        box1_x2 = boxes_preds[..., 2:3]
+        box1_y2 = boxes_preds[..., 3:4]
+        box2_x1 = boxes_labels[..., 0:1]
+        box2_y1 = boxes_labels[..., 1:2]
+        box2_x2 = boxes_labels[..., 2:3]
+        box2_y2 = boxes_labels[..., 3:4]
+
+    x1 = torch.max(box1_x1, box2_x1)
+    y1 = torch.max(box1_y1, box2_y1)
+    x2 = torch.min(box1_x2, box2_x2)
+    y2 = torch.min(box1_y2, box2_y2)
+
+    intersection = (x2 - x1).clamp(0) * (y2 - y1).clamp(0)
+    box1_area = abs((box1_x2 - box1_x1) * (box1_y2 - box1_y1))
+    box2_area = abs((box2_x2 - box2_x1) * (box2_y2 - box2_y1))
+
+    return intersection / (box1_area + box2_area - intersection + 1e-6)
+
+
+def predict(
+    model: torch.nn.Module, image: torch.tensor, iou_threshold: float = 1.0, threshold: float = 0.05
+) -> List[List]:
+    """
+    Apply the model_files to the predictions and to postprocessing
+    Args:
+        model: a trained pytorch model_files.
+        image: image as a torch tensor
+        iou_threshold: a threshold for intersection_over_union function
+        threshold: a threshold for bbox probability
+
+    Returns:
+        predicted bboxes
+
+    """
+    # apply model_files. add a dimension to imitate a batch size of 1
+    logits = model(image[None, :])
+    logging.info('predicted')
+
+    # postprocess. In fact, we could remove indexing with idx here, as there is a single image.
+    # But I prefer to keep it so that this code could be easier changed for cases with batch size > 1
+    bboxes: List[List] = [[] for _ in range(1)]
+    for i in range(3):
+        S = logits[i].shape[2]
+        # it could be better to initialize anchors inside the function, but I don't want to do it for every prediction.
+        anchor = anchors[i] * S
+        boxes_scale_i = cells_to_bboxes(logits[i], anchor, s=S, is_preds=True)
+        for idx, (box) in enumerate(boxes_scale_i):
+            bboxes[idx] += box
+    logging.info('Starting nms')
+    nms_boxes = non_max_suppression(
+        bboxes[idx],
+        iou_threshold=iou_threshold,
+        threshold=threshold,
+        box_format='midpoint',
+    )
+
+    return nms_boxes
+
+
+@st.cache
+def get_model():
+
+    model_name = 'model_files/best_model.pth'
+
+    model = Net()
+    model.load_state_dict(torch.load(model_name))
+    model.eval()
+
+    return model

src/model_architecture.py

@@ -0,0 +1,156 @@
+import torch
+import torch.nn as nn
+
+
+class CNNBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, bn_act=True, **kwargs):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, bias=not bn_act, **kwargs)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.leaky = nn.LeakyReLU(0.1)
+        self.use_bn_act = bn_act
+
+    def forward(self, x):
+        if self.use_bn_act:
+            return self.leaky(self.bn(self.conv(x)))
+        else:
+            return self.conv(x)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels, use_residual=True, num_repeats=1):
+        super().__init__()
+        self.layers = nn.ModuleList()
+        for _ in range(num_repeats):
+            self.layers += [
+                nn.Sequential(
+                    CNNBlock(channels, channels // 2, kernel_size=1),
+                    CNNBlock(channels // 2, channels, kernel_size=3, padding=1),
+                )
+            ]
+
+        self.use_residual = use_residual
+        self.num_repeats = num_repeats
+
+    def forward(self, x):
+        for layer in self.layers:
+            if self.use_residual:
+                x = x + layer(x)
+            else:
+                x = layer(x)
+
+        return x
+
+
+class ScalePrediction(nn.Module):
+    def __init__(self, in_channels, num_classes):
+        super().__init__()
+        self.pred = nn.Sequential(
+            CNNBlock(in_channels, 2 * in_channels, kernel_size=3, padding=1),
+            CNNBlock(2 * in_channels, (num_classes + 5) * 3, bn_act=False, kernel_size=1),
+        )
+        self.num_classes = num_classes
+
+    def forward(self, x):
+        return self.pred(x).reshape(x.shape[0], 3, self.num_classes + 5, x.shape[2], x.shape[3]).permute(0, 1, 3, 4, 2)
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.num_classes = 12
+        self.in_channels = 3
+        # self.config = cfg.model_files.params.config
+        # self.config = [i if i[0] != '(' else literal_eval(i) for i in self.config]
+        self.config = [
+            (32, 3, 1),
+            (64, 3, 2),
+            ['B', 1],
+            (128, 3, 2),
+            ['B', 2],
+            (256, 3, 2),
+            ['B', 8],
+            (512, 3, 2),
+            ['B', 8],
+            (1024, 3, 2),
+            ['B', 4],
+            (512, 1, 1),
+            (1024, 3, 1),
+            'S',
+            (256, 1, 1),
+            'U',
+            (256, 1, 1),
+            (512, 3, 1),
+            'S',
+            (128, 1, 1),
+            'U',
+            (128, 1, 1),
+            (256, 3, 1),
+            'S',
+        ]
+        # print('self.config', self.config)
+        self.layers = self._create_conv_layers()
+
+    def forward(self, x):
+        outputs = []  # for each scale
+        route_connections = []
+        for layer in self.layers:
+            if isinstance(layer, ScalePrediction):
+                outputs.append(layer(x))
+                continue
+            # print(layer, x.shape)
+            x = layer(x)
+
+            if isinstance(layer, ResidualBlock) and layer.num_repeats == 8:
+                route_connections.append(x)
+
+            elif isinstance(layer, nn.Upsample):
+                x = torch.cat([x, route_connections[-1]], dim=1)
+                route_connections.pop()
+
+        return outputs
+
+    def _create_conv_layers(self):
+        layers = nn.ModuleList()
+        in_channels = self.in_channels
+
+        for module in self.config:
+            # print(module, type(module))
+            if isinstance(module, tuple):
+                out_channels, kernel_size, stride = module
+                layers.append(
+                    CNNBlock(
+                        in_channels,
+                        out_channels,
+                        kernel_size=kernel_size,
+                        stride=stride,
+                        padding=1 if kernel_size == 3 else 0,
+                    )
+                )
+                in_channels = out_channels
+
+            elif isinstance(module, list):
+                num_repeats = module[1]
+                layers.append(
+                    ResidualBlock(
+                        in_channels,
+                        num_repeats=num_repeats,
+                    )
+                )
+
+            elif isinstance(module, str):
+                if module == 'S':
+                    layers += [
+                        ResidualBlock(in_channels, use_residual=False, num_repeats=1),
+                        CNNBlock(in_channels, in_channels // 2, kernel_size=1),
+                        ScalePrediction(in_channels // 2, num_classes=self.num_classes),
+                    ]
+                    in_channels = in_channels // 2
+
+                elif module == 'U':
+                    layers.append(
+                        nn.Upsample(scale_factor=2),
+                    )
+                    in_channels = in_channels * 3
+
+        return layers

src/utils.py

@@ -0,0 +1,61 @@
+from typing import List, Dict
+
+import matplotlib
+import matplotlib.patches as patches
+import matplotlib.pyplot as plt
+import numpy as np
+import tomli as tomllib
+
+
+def plot_img_with_rects(
+    img: np.array, boxes: List[List], threshold: float = 0.5, coef: int = 400
+) -> matplotlib.figure.Figure:
+    """
+    Plot image with rectangles.
+
+    Args:
+        img: image as a numpy array
+        boxes: the list of the bboxes
+        threshold: threshold for bbox probability
+        coef: coefficient to multiply images. Can be changed when the original image is a different size
+
+    Returns:
+        image with bboxes
+    """
+    fig, ax = plt.subplots(1, figsize=(4, 4))
+
+    # Display the image
+    ax.imshow(img)
+
+    # Create a Rectangle patch
+    for _, rect in enumerate([b for b in boxes if b[1] > threshold]):
+        label, _, xc, yc, w, h = rect
+        xc, yc, w, h = xc * coef, yc * coef, w * coef, h * coef
+        # the coordinates from center-based to left top corner
+        x = xc - w / 2
+        y = yc - h / 2
+        label = int(label)
+        label = label if label != 10 else 'penis'
+        label = label if label != 11 else 'junk'
+        rect = [x, y, x + w, y + h]
+
+        rect_ = patches.Rectangle(
+            (rect[0], rect[1]), rect[2] - rect[0], rect[3] - rect[1], linewidth=2, edgecolor='blue', facecolor='none'
+        )
+        plt.text(rect[2], rect[1], f'{label}', color='blue')
+        # Add the patch to the Axes
+        ax.add_patch(rect_)
+    return fig
+
+
+def get_config() -> Dict:
+    """
+    Get dict from config.
+
+    Returns:
+        config
+    """
+    with open('config.toml', 'rb') as f:
+        config = tomllib.load(f)
+
+    return config

st_app.py

@@ -0,0 +1,46 @@
+import logging
+
+import numpy as np
+import streamlit as st
+import tomli as tomllib
+from PIL import Image
+from streamlit_drawable_canvas import st_canvas
+
+from src.ml_utils import predict, get_model, transforms
+from src.utils import plot_img_with_rects, get_config
+
+logging.info('Starting')
+
+col1, col2 = st.columns(2)
+
+with col1:
+    # Create a canvas component
+    canvas_result = st_canvas(
+        fill_color='#fff',
+        stroke_width=5,
+        stroke_color='#000',
+        background_color='#fff',
+        update_streamlit=True,
+        height=400,
+        width=400,
+        drawing_mode='freedraw',
+        key='canvas',
+    )
+with col2:
+    data = get_config()
+    logging.info('canvas ready')
+    if canvas_result.image_data is not None:
+        # convert a drawn image into numpy array with RGB from a canvas image with RGBA
+        img = np.array(Image.fromarray(np.uint8(canvas_result.image_data)).convert('RGB'))
+        image = transforms(image=img)['image']
+        logging.info('image augmented')
+        model = get_model()
+        logging.info('model ready')
+        pred = predict(model, image)
+        logging.info('prediction done')
+        threshold = st.slider('Bbox probability slider', min_value=0.0, max_value=1.0, value=0.5)
+
+        fig = plot_img_with_rects(image.permute(1, 2, 0).numpy(), pred, threshold, coef=192)
+        fig.savefig('figure_name1.png')
+        image = Image.open('figure_name1.png')
+        st.image(image)