Merge pull request #1 from Danila-Pechenev/app

.github/workflows/cd.yml

@@ -0,0 +1,19 @@
+name: CD
+
+
+on:
+  push:
+    branches: [main]
+
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: git push --force https://dpechenev:[email protected]/spaces/dpechenev/low_light_image_enhancement main

.github/workflows/ci.yml

@@ -0,0 +1,48 @@
+name: CI
+
+on: [push, pull_request]
+
+jobs:
+  lint:
+    name: Linter
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Run linter
+        uses: psf/black@stable
+        with:
+          options: "--check --verbose -l 120"
+          src: "./app"
+          version: "~= 22.0"
+  mypy:
+    name: Mypy
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+      - name: Install dependencies
+        run: |
+          pip install mypy
+          pip install -r app/requirements.txt
+          python3 -m pip install types-Pillow
+      - name: Run mypy
+        run: mypy .
+  test:
+    name: Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 60
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+      - name: Install dependencies
+        run: |
+          pip install pytest
+          pip install -r app/requirements.txt
+      - name: Run tests
+        run: pytest

.gitignore

@@ -127,3 +127,5 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+.idea/

README.md

@@ -1,2 +1,5 @@
 # low-light-image-enhancement
 Mini-application for low-light image enhancement based on keras-io/lowlight-enhance-mirnet model & Streamlit
+
+## Description
+Photos are often taken in poor light. This tool will help you lighten the image using a neural network.

app/app.py

@@ -0,0 +1,50 @@
+import streamlit as st
+from tensorflow import keras
+from PIL import Image
+import io
+
+import model
+
+
+def configure():
+    st.set_page_config(page_title="Low-light image enhancement")
+    if "model" not in st.session_state:
+        st.session_state["model"]: keras.Model = model.create_model()
+
+
+def describe_service():
+    st.title("Low-light image enhancement")
+    st.subheader("Just upload your low-light image and get the processed one!")
+
+
+@st.experimental_memo
+def call_model(uploaded_file: io.BytesIO) -> Image.Image:
+    return model.run_model(uploaded_file, st.session_state["model"])
+
+
+def process_image():
+    uploaded_file = st.file_uploader(
+        label="Choose a file (you can upload new files without refreshing the page)",
+        type=["png", "jpg", "jpeg"],
+    )
+    if uploaded_file:
+        placeholder = st.empty()
+        placeholder.info("The image is being processed. It may take some time. Wait, please...")
+        image = call_model(uploaded_file)
+        placeholder.empty()
+        placeholder.image(image)
+        image_bytes = io.BytesIO()
+        image.save(image_bytes, format="png")
+        st.download_button(
+            label="Download lightened image", data=image_bytes, file_name="lightened.png", mime="image/png"
+        )
+
+
+def main():
+    describe_service()
+    process_image()
+
+
+if __name__ == "__main__":
+    configure()
+    main()

app/model.py

@@ -0,0 +1,29 @@
+from huggingface_hub import from_pretrained_keras
+from tensorflow import keras
+from PIL import Image
+import numpy as np
+import io
+import os
+
+
+def create_model() -> keras.Model:
+    return from_pretrained_keras("keras-io/lowlight-enhance-mirnet")
+
+
+def run_model(image_bytes: io.BytesIO, model: keras.Model) -> Image.Image:
+    image = Image.open(image_bytes)
+    width, height = image.size
+    image = image.resize((960, 640))
+    image_array = np.expand_dims(keras.utils.img_to_array(image).astype("float32") / 255.0, axis=0)
+    output = model.predict(image_array)
+    output_image_array = (output[0] * 255.0).clip(0, 255)
+    output_image_array = output_image_array.reshape(
+        (np.shape(output_image_array)[0], np.shape(output_image_array)[1], 3)
+    ).astype(np.uint8)
+    output_image = Image.fromarray(output_image_array).resize((width, height))
+
+    if not os.path.exists("user_data"):
+        os.makedirs("user_data")
+    output_image.save("user_data/output.jpg")
+
+    return output_image

app/requirements.txt

@@ -0,0 +1,4 @@
+numpy==1.23.5
+streamlit==1.15.2
+tensorflow==2.11.0
+huggingface-hub==0.11.1

setup.cfg

@@ -0,0 +1,2 @@
+[mypy]
+ignore_missing_imports = True

telegram_bot/bot.py

@@ -0,0 +1,47 @@
+import telebot
+from PIL import Image
+import sys
+import os
+import io
+
+sys.path.append(os.path.join(os.path.dirname(__file__), "..", "app"))
+import model
+
+bot = telebot.TeleBot("<Paste your token here>")
+model_instance = model.create_model()
+
+
+@bot.message_handler(content_types=["photo"])
+def photo(message):
+    user_id = message.from_user.id
+    file_id = message.photo[-1].file_id
+    file_info = bot.get_file(file_id)
+    downloaded_file = bot.download_file(file_info.file_path)
+
+    if not os.path.exists("user_data"):
+        os.makedirs("user_data")
+    filename = f"user_data/image_{user_id}_{file_id}.jpg"
+    with open(filename, "wb") as new_file:
+        new_file.write(downloaded_file)
+
+    bot.send_message(message.from_user.id, "Изображение обрабатывается...")
+
+    image = Image.open(filename)
+    image_bytes = io.BytesIO()
+    image.save(image_bytes, format=image.format)
+    output_image = model.run_model(image_bytes, model_instance)
+    os.remove(filename)
+
+    bot.send_photo(user_id, output_image)
+
+
+@bot.message_handler(content_types=["text"])
+def start(message):
+    bot.send_message(
+        message.from_user.id,
+        """Этот бот осветляет фотографии, сделанные при плохом освещении, при помощи нейронной сети.\
+        Просто пришлите сюда фото!""",
+    )
+
+
+bot.polling(none_stop=True, interval=0)

test/test_model.py

@@ -0,0 +1,36 @@
+from PIL import Image
+import io
+import os
+from ..app import model
+
+model_instance = model.create_model()
+
+
+def template(filename: str):
+    image = Image.open(filename)
+    width, height = image.size
+    image_bytes = io.BytesIO()
+    image.save(image_bytes, format=image.format)
+    output_image = model.run_model(image_bytes, model_instance)
+    return width, height, output_image
+
+
+def test_image_jpg():
+    width, height, output_image = template(os.path.join(os.getcwd(), "test/test_images/test1.jpg"))
+
+    assert width == output_image.size[0]
+    assert height == output_image.size[1]
+
+
+def test_image_png():
+    width, height, output_image = template(os.path.join(os.getcwd(), "test/test_images/test2.png"))
+
+    assert width == output_image.size[0]
+    assert height == output_image.size[1]
+
+
+def test_image_jpeg():
+    width, height, output_image = template(os.path.join(os.getcwd(), "test/test_images/test3.jpeg"))
+
+    assert width == output_image.size[0]
+    assert height == output_image.size[1]