init commit

.gitignore

@@ -0,0 +1,5 @@
+/venv
+__pycache__
+/test_*
+/.idea
+.env

README.md

@@ -1,6 +1,6 @@
 ---
-title: Mist Checker
-emoji: 🌖
+title: Shit Checker
+emoji: 💩
 colorFrom: blue
 colorTo: indigo
 sdk: gradio

app.py

@@ -0,0 +1,47 @@
+import os
+import textwrap
+import time
+
+import gradio as gr
+from PIL import Image
+
+from check import DEFAULT_MODEL, predict, MODELS
+
+
+def _predict_fn(image: Image.Image, model_name: str = DEFAULT_MODEL, max_batch_size: int = 8):
+    start_time = time.time()
+    result = predict(image, model_name, max_batch_size)
+    duration = time.time() - start_time
+    info = f'Time cost: **{duration:.3f}s**'
+    return result, info
+
+
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr_info = gr.Markdown(textwrap.dedent("""
+                 Quickly check if the image is **glazed or misted** (we call it shat💩).
+
+                 And then you can just remove these shit with
+                 [mf666/mist-fucker](https://huggingface.co/spaces/mf666/mist-fucker),
+                 without fucking the normal images (no detail losses).
+            """).strip())
+
+        with gr.Row():
+            with gr.Column():
+                gr_input = gr.Image(label='Image To Check', image_mode='RGB', type='pil')
+                gr_models = gr.Dropdown(choices=MODELS, value=DEFAULT_MODEL, label='Models')
+                gr_max_batch_size = gr.Slider(minimum=1, maximum=16, value=8, step=1, label='Max Batch Size')
+                gr_submit = gr.Button(value='Check The Shit', variant='primary')
+
+            with gr.Column():
+                gr_label = gr.Label(label='Check Result')
+                gr_time_cost = gr.Markdown(label='Information')
+
+            gr_submit.click(
+                _predict_fn,
+                inputs=[gr_input, gr_models, gr_max_batch_size],
+                outputs=[gr_label, gr_time_cost],
+            )
+
+    demo.queue(os.cpu_count()).launch()

check.py

@@ -0,0 +1,81 @@
+import math
+import os
+import random
+from functools import lru_cache
+
+import numpy as np
+from PIL import Image
+from huggingface_hub import HfFileSystem, HfApi
+from imgutils.utils import open_onnx_model
+from natsort import natsorted
+
+hf_token = os.environ.get('HF_TOKEN')
+hf_fs = HfFileSystem(token=hf_token)
+hf_client = HfApi(token=hf_token)
+
+REPOSITORY = 'mf666/shit-checker'
+MODELS = natsorted([
+    os.path.splitext(os.path.relpath(file, REPOSITORY))[0]
+    for file in hf_fs.glob(f'{REPOSITORY}/*.onnx')
+])
+DEFAULT_MODEL = 'mobilenet.small'
+
+
+@lru_cache()
+def _open_model(model_name):
+    return open_onnx_model(hf_client.hf_hub_download(REPOSITORY, f'{model_name}.onnx'))
+
+
+_DEFAULT_ORDER = 'HWC'
+
+
+def _get_hwc_map(order_):
+    return tuple(_DEFAULT_ORDER.index(c) for c in order_.upper())
+
+
+def _encode_channels(image, channels_order='CHW', is_float=True):
+    array = np.asarray(image.convert('RGB'))
+    array = np.transpose(array, _get_hwc_map(channels_order))
+    if not is_float:
+        assert array.dtype == np.uint8
+    else:
+        array = (array / 255.0).astype(np.float32)
+        assert array.dtype == np.float32
+    return array
+
+
+def _img_encode(image, size=(384, 384), normalize=(0.5, 0.5)):
+    image = image.resize(size, Image.BILINEAR)
+    data = _encode_channels(image, channels_order='CHW')
+
+    if normalize is not None:
+        mean_, std_ = normalize
+        mean = np.asarray([mean_]).reshape((-1, 1, 1))
+        std = np.asarray([std_]).reshape((-1, 1, 1))
+        data = (data - mean) / std
+
+    return data.astype(np.float32)
+
+
+def _raw_predict(images, model_name=DEFAULT_MODEL):
+    items = []
+    for image in images:
+        items.append(_img_encode(image.convert('RGB')))
+    input_ = np.stack(items)
+    output, = _open_model(model_name).run(['output'], {'input': input_})
+    return output.mean(axis=0)
+
+
+def predict(image, model_name=DEFAULT_MODEL, max_batch_size=8):
+    area = image.width * image.height
+    batch_size = int(max(min(math.ceil(area / (384 * 384)) + 1, max_batch_size), 1))
+    blocks = []
+    for _ in range(batch_size):
+        x0 = random.randint(0, max(0, image.width - 384))
+        y0 = random.randint(0, max(0, image.height - 384))
+        x1 = min(x0 + 384, image.width)
+        y1 = min(y0 + 384, image.height)
+        blocks.append(image.crop((x0, y0, x1, y1)))
+
+    scores = _raw_predict(blocks, model_name)
+    return dict(zip(['shat', 'normal'], map(lambda x: x.item(), scores)))

requirements.txt

@@ -0,0 +1,12 @@
+gradio==4.12.0
+numpy
+pillow
+onnxruntime
+huggingface_hub>=0.14.0
+scikit-image
+pandas
+opencv-python>=4.6.0
+hbutils>=0.9.0
+dghs-imgutils>=0.1.0
+httpx==0.23.0
+natsort