artcracker commited on
Commit
c086220
·
verified ·
1 Parent(s): dcbbf7f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +100 -0
app.py ADDED
@@ -0,0 +1,100 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoImageProcessor, AutoModel
3
+ import torch
4
+ from PIL import Image
5
+ import os
6
+ import json
7
+ import uuid
8
+ import numpy as np
9
+
10
+ # Загрузка модели DINOv2
11
+ processor = AutoImageProcessor.from_pretrained("facebook/dino-vits16")
12
+ model = AutoModel.from_pretrained("facebook/dino-vits16")
13
+
14
+ # Функция для извлечения эмбеддинга
15
+ def extract_features(img):
16
+ inputs = processor(images=img, return_tensors="pt")
17
+ with torch.no_grad():
18
+ outputs = model(**inputs)
19
+ return outputs.last_hidden_state[:, 0].squeeze().numpy()
20
+
21
+ # Косинусное расстояние между двумя эмбеддингами
22
+ def cosine_similarity(vec1, vec2):
23
+ a = np.array(vec1)
24
+ b = np.array(vec2)
25
+ return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
26
+
27
+ # Путь к базе
28
+ db_path = "embeddings.json"
29
+
30
+ # Сохранение фото в базу
31
+ def save_embedding(image, building_name):
32
+ if not building_name:
33
+ return "Ошибка: нужно указать название здания."
34
+
35
+ embedding = extract_features(image).tolist()
36
+ entry = {
37
+ "id": str(uuid.uuid4()),
38
+ "building_name": building_name,
39
+ "embedding": embedding
40
+ }
41
+
42
+ if os.path.exists(db_path):
43
+ with open(db_path, "r") as f:
44
+ data = json.load(f)
45
+ else:
46
+ data = []
47
+
48
+ data.append(entry)
49
+
50
+ with open(db_path, "w") as f:
51
+ json.dump(data, f, indent=2)
52
+
53
+ return f"Фото сохранено для здания: {building_name}"
54
+
55
+ # Поиск похожего здания
56
+ def identify_building(image):
57
+ if not os.path.exists(db_path):
58
+ return "База данных пуста. Добавьте здания."
59
+
60
+ with open(db_path, "r") as f:
61
+ data = json.load(f)
62
+
63
+ if not data:
64
+ return "База данных пуста."
65
+
66
+ embedding = extract_features(image).tolist()
67
+
68
+ similarities = []
69
+ for item in data:
70
+ score = cosine_similarity(item["embedding"], embedding)
71
+ similarities.append({
72
+ "building_name": item["building_name"],
73
+ "score": score
74
+ })
75
+
76
+ similarities.sort(key=lambda x: x["score"], reverse=True)
77
+ best = similarities[0]
78
+
79
+ return f"Похоже на: {best['building_name']}\nСовпадение: {best['score']:.4f}"
80
+
81
+ # Интерфейс
82
+ with gr.Blocks() as demo:
83
+ gr.Markdown("## 🏛 Распознавание зданий и пополнение базы")
84
+
85
+ with gr.Tab("🔍 Найти здание"):
86
+ with gr.Row():
87
+ img_input = gr.Image(type="pil")
88
+ recognize_button = gr.Button("Распознать здание")
89
+ result_text = gr.Textbox()
90
+ recognize_button.click(fn=identify_building, inputs=img_input, outputs=result_text)
91
+
92
+ with gr.Tab("➕ Добавить новое здание"):
93
+ with gr.Row():
94
+ img_save = gr.Image(type="pil")
95
+ building_name = gr.Textbox(label="Название здания")
96
+ save_button = gr.Button("Сохранить в базу")
97
+ save_result = gr.Textbox()
98
+ save_button.click(fn=save_embedding, inputs=[img_save, building_name], outputs=save_result)
99
+
100
+ demo.launch()