Merge pull request #2 from timurgainutdinov1/module_4_1

Module 4 1

image_classification_streamlit.py

@@ -2,10 +2,26 @@
 
 import requests
 import streamlit as st
-from PIL import Image
+import translators as ts
+from PIL import Image, UnidentifiedImageError
+from requests.exceptions import MissingSchema
 from transformers import ViTForImageClassification, ViTImageProcessor
 
 
+class MissingSourceError(Exception):
+    """Класс представляет ошибку,
+    возникающую при отсутствии
+    источника изображения."""
+    pass
+
+
+class TwoSourcesError(Exception):
+    """Класс представляет ошибку,
+    возникающую при указании
+    двух источников изображений."""
+    pass
+
+
 @st.cache_resource
 def load_model():
     """Загрузка модели"""
@@ -25,11 +41,21 @@ def get_image_link():
     return st.text_input("Введите ссылку на изображение для распознавания")
 
 
-def load_image(url):
+def get_image_file():
+    """Загрузка файла с изображением."""
+    return st.file_uploader("Или загрузите изображение из файла")
+
+
+def load_image_from_url(url):
     """Загрузка изображения из указанного URL-адреса
     с помощью библиотеки requests."""
     img = Image.open(requests.get(url, stream=True).raw)
-    st.image(img)
+    return img
+
+
+def load_image_from_file(file):
+    """Загрузка изображения из файла."""
+    img = Image.open(file)
     return img
 
 
@@ -57,20 +83,56 @@ def show_results(results):
 
 st.title("Модель для классификации изображений vit-base-patch16-224")
 
-link = get_image_link()
+image_link = get_image_link()
+image_file = get_image_file()
 
 result = st.button("Распознать изображение")
 
 if result:
     try:
-        loaded_image = load_image(link)
+        loaded_image = ""
+        if image_link != "" and image_file is not None:
+            raise TwoSourcesError
+        elif image_link != "":
+            loaded_image = load_image_from_url(image_link)
+        elif image_file is not None:
+            loaded_image = load_image_from_file(image_file)
+        else:
+            raise MissingSourceError
+        st.image(loaded_image)
         with st.spinner("Идет обработка... Пожалуйста, подождите..."):
             result = image_classification(loaded_image)
-        st.markdown(f"Результаты распознавания: :rainbow[{result}]")
-        st.snow()
-    except IOError:
+        translated_result = ts.translate_text(result,
+                                              translator="bing",
+                                              from_language="en",
+                                              to_language="ru")
+        st.markdown(f"Результаты распознавания: {translated_result}")
+    except MissingSourceError:
+        st.error(
+            "Вы не предоставили источник "
+            "для загрузки изображения. "
+            "Загрузите файл с изображением или укажите ссылку "
+            "и попробуйте снова!",
+            icon="😞",
+        )
+    except MissingSchema:
+        st.error(
+            "Некорректная ссылка! "
+            "Укажите корректную ссылку "
+            "и попробуйте снова!",
+            icon="😞",
+        )
+    except UnidentifiedImageError:
+        st.error(
+            "Ваша ссылка или файл не содержат изображения. "
+            "Предоставьте корректную ссылку или файл "
+            "и попробуйте снова!",
+            icon="😞",
+        )
+    except TwoSourcesError:
         st.error(
-            "Не удалось найти изображение по указанной ссылке. "
-            "Попробуйте снова!",
+            "Вы указали два источника. "
+            "Удалите один из источников "
+            "и попробуйте снова!",
             icon="😞",
         )

requirements.txt

@@ -4,4 +4,5 @@ transformers[torch]
 requests
 Pillow
 pydantic
-streamlit
+streamlit
+translators

test_app.py

@@ -8,14 +8,15 @@
                        default_timeout=1000).run()
 
 
-def test_incorrect_url():
+def test_no_image_url():
     """Проверка ввода URL-адреса на объект,
-    который не является изображением."""
+    который не является изображением"""
     at.text_input[0].set_value("https://www.google.com/").run()
     at.button[0].click().run()
     assert at.error[0].value == (
-        "Не удалось найти изображение по указанной ссылке. "
-        "Попробуйте снова!"
+        "Ваша ссылка или файл не содержат изображения. "
+        "Предоставьте корректную ссылку или файл "
+        "и попробуйте снова!"
     )
 
 
@@ -24,8 +25,10 @@ def test_null_url():
     at.text_input[0].set_value("").run()
     at.button[0].click().run()
     assert at.error[0].value == (
-        "Не удалось найти изображение по указанной ссылке. "
-        "Попробуйте снова!"
+        "Вы не предоставили источник "
+        "для загрузки изображения. "
+        "Загрузите файл с изображением или укажите ссылку "
+        "и попробуйте снова!"
     )
 
 
@@ -43,6 +46,16 @@ def test_correct_url():
     at.button[0].click().run()
     time.sleep(5)
     assert at.markdown[0].value == (
-        "Результаты распознавания: "
-        ":rainbow[tabby, tabby cat]"
+        "Результаты распознавания: табби, полосатый кот"
+    )
+
+
+def test_incorrect_url():
+    """Проверка ввода некорректного URL-адреса"""
+    at.text_input[0].set_value("1234").run()
+    at.button[0].click().run()
+    assert at.error[0].value == (
+        "Некорректная ссылка! "
+        "Укажите корректную ссылку "
+        "и попробуйте снова!"
     )