hw4: done task 2

hw4/.gitignore

@@ -3,6 +3,10 @@ venv
 .idea
 .ipynb_checkpoints
 task1_input/*
-!task1_input/.gitkeep
 task2_input/*
+task3_input/*
+task4_input/*
+!task1_input/.gitkeep
 !task2_input/.gitkeep
+!task3_input/.gitkeep
+!task4_input/.gitkeep

hw4/task2.py

-import cv2 as cv
+import json
+from pathlib import Path
+from typing import List
+
 import numpy as np
+import streamlit as st
+from PIL import Image
+from deepface import DeepFace
+from faiss import IndexFlatL2
+
+
+def get_image_embeddings(image: np.ndarray):
+    return DeepFace.represent(
+        img_path=image,
+        model_name="Facenet",
+        max_faces=1,
+    )
+
+
+def create_dataset_embeddings() -> (np.ndarray, List):
+    staff_dataset = np.loadtxt(
+        fname="task2_input/staff_photo.csv",
+        delimiter=",",
+        dtype=str,
+    )
+    embeddings_list = []
+    ids = []
+    index = 0
+
+    for path in Path("task2_input/hse_faces_miem/").rglob("*.jpeg"):
+        path = str(path)
+        print(path)
+        image = np.array(Image.open(path))[:, :, ::-1]
+
+        try:
+            detections = get_image_embeddings(image)
+        except Exception:
+            continue
+
+        if len(detections):
+            embeddings = detections[0]["embedding"]
+            id = int(path.replace("/", "_").split("_")[-2])
+            name = staff_dataset[
+                np.where(staff_dataset[:, 0] == str(id))[0][0]
+            ][1]
+
+            embeddings_list.append(embeddings)
+            ids.append(
+                {
+                    "index": index,
+                    "id": id,
+                    "name": name,
+                    "path": path,
+                }
+            )
+
+            index += 1
+    return np.array(embeddings_list), ids
+
+
+def dump_embeddings(vectors, ids):
+    np.savetxt("embeddings.txt", vectors)
+    with open("ids.txt", "w") as file:
+        json.dump(ids, file)
+
+
+def load_embeddings():
+    vectors = np.loadtxt("embeddings.txt")
+    with open("ids.txt") as f:
+        ids = json.load(f)
+    return vectors, ids
+
+
+def main(index, ids):
+
+    st.title("MIEM Lookalike")
+
+    column1, column2 = st.columns(2)
+
+    with column1:
+        st.markdown("### Ваше фото")
+        preview_placeholder = st.empty()
+
+        uploaded_file = st.file_uploader(
+            "Выберите изображение",
+            type=["jpg", "jpeg", "png"],
+        )
+
+        if uploaded_file is not None:
+            preview_placeholder.image(
+                Image.open(uploaded_file),
+                caption="Выбранное изображение",
+                use_container_width=True,
+            )
+
+    with column2:
+        st.markdown("### Результат")
+        result_image_placeholder = st.empty()
+        result_text_placeholder = st.empty()
+
+        if uploaded_file is not None:
+
+            np_image = np.asarray(
+                Image.open(uploaded_file)
+            )[:, :, ::-1]
+
+            detection = get_image_embeddings(np_image)[0]
+
+            distances, indices = index.search(
+                np.array([detection["embedding"]]),
+                1,
+            )
+            print(distances)
+
+            detected_id = indices[0][0]
+            detected_name = str(ids[detected_id]["name"])
+            path = str(ids[detected_id]["path"])
+
+            result_image_placeholder.image(
+                Image.open(path),
+                caption="Похожее изображение",
+                use_container_width=True,
+            )
+
+            result_text_placeholder.text_area(
+                f"Результат:",
+                detected_name,
+            )
+
+    st.markdown("""
+        <style>
+            .stButton>button {
+                width: 100%;
+            }
+            .upload-text {
+                text-align: center;
+                padding: 20px;
+            }
+        </style>
+        """,
+        unsafe_allow_html=True,
+    )
+
 
+if __name__ == "__main__":
+    # dump_embeddings(*create_dataset_embeddings())
 
+    embeddings, ids = load_embeddings()
+    index = IndexFlatL2(128)
+    index.add(embeddings)
+    main(index, ids)