Add main.py, change model_execute.py

Former-commit-id: 4da2a83f1b275ec8320c2ed7bed56916c863e44b
Former-commit-id: 962450d9b097706e9622530cff0fa3091aa8a348

.gitignore

@@ -1 +1 @@
-.idea/
+__pycache__

main.py

@@ -0,0 +1,91 @@
+import streamlit as st
+import torch
+import pickle
+from PIL import Image
+import io
+
+from model_execute import preprocess_images, output_to_names
+from summarization import init_model_and_tokenizer, summarize
+from wikipedia_api import getWikipedia
+
+@st.cache_resource
+def load_recognition_model():
+    """
+    Loads the translation model pipeline.
+    """
+    filename = "pickle_model.pkl"
+    with open(filename, 'rb') as file:
+        model = pickle.load(file)
+    return model
+
+@st.cache_resource
+def load_summarizer():
+    """
+    Loads the summarization model.
+    """
+
+    summarizer, tokenizer = init_model_and_tokenizer()
+    return summarizer, tokenizer
+
+def predict_images(images, model):
+    """
+    Predicts each landmark name in `images` list.
+    """
+    images = preprocess_images(images)
+
+    with torch.no_grad():
+        output = model(images)
+    
+    names = output_to_names(output)
+
+    return names
+
+def load_images():
+    """
+    Loads user's images.
+    """
+    uploaded_files = st.file_uploader(
+            label="Загрузите ваши фотографии.",
+            type=['png', 'jpg'],
+            accept_multiple_files=True
+        )
+    if uploaded_files is not None:
+        images = []
+        for file in uploaded_files:
+            image_data = file.getvalue()
+            st.image(image_data)
+            images.append(image_data)
+        
+        return [Image.open(io.BytesIO(image_data)) for image_data in images]
+    else:
+        return None
+
+
+# Load models
+landmark_model = load_recognition_model()
+summarizer, tokenizer = load_summarizer()
+
+st.title("Распознавание достопримечательностей")
+
+# Images input.
+images = load_images()
+
+result = st.button('Распознать')
+
+if result:
+    # Get predictions
+    names = predict_images(images, landmark_model)
+    
+    st.write(names)
+    # Request descriptions and coordinates from Wikipedia.
+    wiki_data = getWikipedia(names)
+
+    # Show summarized descriptions for each landmark.
+    for landmark in wiki_data:
+        description = landmark['summary']
+        summarized = summarize(description, summarizer, tokenizer)
+        landmark['summarized'] = summarized
+
+    st.write(wiki_data)
+
+    # Draw a map.

model_execute.py

@@ -2,7 +2,36 @@ import pickle
 import torch
 import torchvision.transforms as transforms
 from PIL import Image
-import torchvision
+import csv
+
+def preprocess_images(images):
+    """
+    Preprocess image for the model.
+    """
+    preprocess = transforms.Compose([
+        transforms.Resize([70, 70]),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    images_tensor = [preprocess(image) for image in images]
+    image_batch = torch.stack(images_tensor)
+
+    return image_batch
+
+def output_to_names(output):
+    """
+    Converts model outputs to category names names.
+    """
+    with open('cat.csv') as file:
+        reader = csv.reader(file)
+        cat_list = list(reader)[0]
+
+    names = []
+    for prediction in output:
+        probabilities = torch.nn.functional.softmax(prediction, dim=0)
+        index = probabilities.argmax()
+        names.append(cat_list[index])
+    return names
 
 def check_photo(name, photo):
     preprocess = transforms.Compose([
@@ -27,39 +56,23 @@ def check_photo(name, photo):
     probabilities = torch.nn.functional.softmax(output[0], dim=0)
     print(name, probabilities)
 
-pkl_filename = "pickle_model.pkl"
-with open(pkl_filename, 'rb') as file:
-    model = pickle.load(file)
 
-model.eval()
-# sample execution (requires torchvision)
+if __name__ == "__main__":
+    pkl_filename = "pickle_model.pkl"
+    with open(pkl_filename, 'rb') as file:
+        model = pickle.load(file)
+
+    model.eval()
+    # sample execution (requires torchvision)
 
-gates_photo = Image.open("gates500.jpg")
-musk_photo = Image.open("mask.jpg")
-bezos_photo = Image.open("bezos500.jpg")
-zuker_photo = Image.open("zuckerberg500.jpg")
-jobs_photo = Image.open("jobs500.jpg")
-test_photos_dict = {'gates':gates_photo, 'musk':musk_photo, 'bezos':bezos_photo,'zuker': zuker_photo,'jobs': jobs_photo}
-for name in test_photos_dict:
-    check_photo(name, test_photos_dict[name])
-# preprocess = transforms.Compose([
-#     transforms.Resize(256),
-#     transforms.CenterCrop(224),
-#     transforms.ToTensor(),
-#     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-# ])
-# input_tensor = preprocess(test_photos_list)
-# input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
-#
-# # move the input and model to GPU for speed if available
-# if torch.cuda.is_available():
-#     input_batch = input_batch.to('cuda')
-#     model.to('cuda')
-#
-# with torch.no_grad():
-#     output = model(input_batch)
-# # Tensor of shape 1000, with confidence scores over ImageNet's 1000 classes
-# print(output[0])
-# print(model)
+    gates_photo = Image.open("gates500.jpg")
+    musk_photo = Image.open("mask.jpg")
+    bezos_photo = Image.open("bezos500.jpg")
+    zuker_photo = Image.open("zuckerberg500.jpg")
+    jobs_photo = Image.open("jobs500.jpg")
+    test_photos_dict = {'gates':gates_photo, 'musk':musk_photo, 'bezos':bezos_photo,'zuker': zuker_photo,'jobs': jobs_photo}
+    for name in test_photos_dict:
+        check_photo(name, test_photos_dict[name])
 
-# print(probabilities)
+    tensor = torch.tensor([[-1.8637, -1.6411, -1.5038, -2.9645, -1.8477, 6.5004], [-1.6067, -1.6597, -1.0925, 5.1295, -1.6491, -1.4739], [-0.2427, -0.6140, -1.1936, -2.1147, 4.8429, -2.0129]])
+    print(output_to_names(tensor))

requirements.txt

@@ -5,4 +5,6 @@ bs4
 lxml
 requests
 requests
-transformers
+transformers
+streamlit
+wikipedia