add random crop

app.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
-import cv2
+
+import random
+
 import gradio as gr
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
 from CCAgT_utils.categories import CategoriesInfos
-from CCAgT_utils.slice import __create_xy_slice
 from CCAgT_utils.types.mask import Mask
 from CCAgT_utils.visualization import plot
 from PIL import Image
@@ -68,42 +69,56 @@ def colorize(
     return mask.colorized(CategoriesInfos()) / 255
 
 
-def check_and_resize(
-    image: np.ndarray,
-) -> np.ndarray:
+# Copied from https://github.com/albumentations-team/albumentations/blob/b1af92ab8e57279f5acd5987770a86a8d6b6b0e5/albumentations/augmentations/crops/functional.py#L35
+def get_random_crop_coords(
+    height: int,
+    width: int,
+    crop_height: int,
+    crop_width: int,
+    h_start: float,
+    w_start: float,
+):
+    y1 = int((height - crop_height + 1) * h_start)
+    y2 = y1 + crop_height
+    x1 = int((width - crop_width + 1) * w_start)
+    x2 = x1 + crop_width
+    return x1, y1, x2, y2
+
 
-    if image.shape[0] > 1200 or image.shape[1] > 1600:
-        r = 1600.0 / image.shape[1]
-        dim = (1600, int(image.shape[0] * r))
-        return cv2.resize(image, dim, interpolation=cv2.INTER_AREA)
+# Copied from https://github.com/albumentations-team/albumentations/blob/b1af92ab8e57279f5acd5987770a86a8d6b6b0e5/albumentations/augmentations/crops/functional.py#L46
 
-    return image
+
+def random_crop(
+    img: np.ndarray,
+    crop_height: int,
+    crop_width: int,
+    h_start: float,
+    w_start: float,
+) -> np.ndarray:
+    height, width = img.shape[:2]
+
+    x1, y1, x2, y2 = get_random_crop_coords(
+        height, width, crop_height, crop_width, h_start, w_start,
+    )
+    img = img[y1:y2, x1:x2]
+    return img
 
 
 def process_big_images(
     image: Image.Image,
 ) -> Mask:
     '''Process and post-processing for images bigger than 400x300'''
-    img = check_and_resize(np.asarray(image))
-    mask = np.zeros(shape=(img.shape[0], img.shape[1]), dtype=np.uint8)
+    img = np.asarray(image)
 
-    for bbox in __create_xy_slice(image.size[1], image.size[0], 300, 400):
-        part = cv2.copyMakeBorder(
-            img,
-            bbox.y_init,
-            bbox.y_end,
-            bbox.x_init,
-            bbox.x_end,
-            cv2.BORDER_REFLECT,
-        )
-        target_size = (part.shape[0], part.shape[1])
+    if img.shape[0] > 300 or img.shape[1] > 400:
+        img = random_crop(img, 300, 400, random.random(), random.random())
 
-        outputs = segment(Image.fromarray(part))
-        msk = post_processing(outputs, target_size)
+    target_size = (img.shape[0], img.shape[1])
 
-        mask[bbox.slice_y, bbox.slice_x] = msk[bbox.slice_y, bbox.slice_x]
+    outputs = segment(Image.fromarray(img))
+    msk = post_processing(outputs, target_size)
 
-    return Mask(mask)
+    return Mask(msk)
 
 
 def image_with_mask(
@@ -156,7 +171,8 @@ This is demo for the SegFormer fine-tuned on sub-dataset from
 [CCAgT dataset](https://huggingface.co/datasets/lapix/CCAgT). This model
 was trained to segment cervical cells silver-stained (AgNOR technique)
 images with resolution of 400x300. The model was available at HF hub at
-[{model_hub_name}](https://huggingface.co/{model_hub_name}).
+[{model_hub_name}](https://huggingface.co/{model_hub_name}). If input
+an image bigger than 400x300, the demo will random crop it.
 """
 examples = [
     [f'https://hf.co/{model_hub_name}/resolve/main/sampleA.png'],

requirements.txt

@@ -1,6 +1,5 @@
 CCAgT-utils
 matplotlib
 numpy
-opencv-python
 torch
 transformers