David Ko
feat(vision-rag): add LangChain deps and verify .venv; vision_rag_query uses ChatOpenAI; set default OPENAI_MODEL to gpt-4o
637dbbd
raw
history blame
66.4 kB
# -*- coding: utf-8 -*-
# Set matplotlib config directory to avoid permission issues
import os
os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib'
from flask import Flask, request, jsonify, send_from_directory, redirect, url_for, session, render_template_string, make_response
from datetime import timedelta
import torch
from PIL import Image
import numpy as np
import io
from io import BytesIO
import base64
import uuid
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import time
from flask_cors import CORS
import json
import sys
import requests
try:
from openai import OpenAI
except Exception as _e:
OpenAI = None
try:
# LangChain for RAG answering
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
except Exception as _e:
ChatOpenAI = None
ChatPromptTemplate = None
StrOutputParser = None
from flask_login import (
LoginManager,
UserMixin,
login_user,
logout_user,
login_required,
current_user,
fresh_login_required,
login_fresh,
)
# Fix for SQLite3 version compatibility with ChromaDB
try:
import pysqlite3
sys.modules['sqlite3'] = pysqlite3
except ImportError:
print("Warning: pysqlite3 not found, using built-in sqlite3")
import chromadb
from chromadb.utils import embedding_functions
app = Flask(__name__, static_folder='static')
app.secret_key = 'your_secret_key_here' # ์„ธ์…˜ ์•”ํ˜ธํ™”๋ฅผ ์œ„ํ•œ ๋น„๋ฐ€ ํ‚ค
app.config['CORS_HEADERS'] = 'Content-Type'
# Remember cookie (Flask-Login) โ€” minimize duration to prevent auto re-login
app.config['REMEMBER_COOKIE_DURATION'] = timedelta(seconds=1)
app.config['REMEMBER_COOKIE_SECURE'] = True # Spaces uses HTTPS
app.config['REMEMBER_COOKIE_HTTPONLY'] = True
app.config['REMEMBER_COOKIE_SAMESITE'] = 'None'
# Session cookie (Flask-Session)
app.config['SESSION_COOKIE_SECURE'] = True # HTTPS
app.config['SESSION_COOKIE_HTTPONLY'] = True
app.config['SESSION_COOKIE_SAMESITE'] = 'None'
app.config['SESSION_COOKIE_PATH'] = '/'
CORS(app) # Enable CORS for all routes
# ์‹œํฌ๋ฆฟ ํ‚ค ์„ค์ • (์„ธ์…˜ ์•”ํ˜ธํ™”์— ์‚ฌ์šฉ)
app.config['SECRET_KEY'] = os.environ.get('SECRET_KEY', 'vision_llm_agent_secret_key')
app.config['SESSION_TYPE'] = 'filesystem'
app.config['PERMANENT_SESSION_LIFETIME'] = timedelta(seconds=120) # ์„ธ์…˜ ์œ ํšจ ์‹œ๊ฐ„ (2๋ถ„)
app.config['SESSION_REFRESH_EACH_REQUEST'] = False # ์ ˆ๋Œ€ ๋งŒ๋ฃŒ(๋กœ๊ทธ์ธ ๊ธฐ์ค€ 2๋ถ„ ํ›„ ๋งŒ๋ฃŒ)
# Flask-Login ์„ค์ •
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'login'
login_manager.session_protection = 'strong'
# When authentication is required or session is not fresh, redirect to login instead of 401
login_manager.refresh_view = 'login'
@login_manager.unauthorized_handler
def handle_unauthorized():
# For non-authenticated access, send user to login
return redirect(url_for('login'))
@login_manager.needs_refresh_handler
def handle_needs_refresh():
# For non-fresh sessions (e.g., after expiry or only remember-cookie), send to login
return redirect(url_for('login'))
# ์„ธ์…˜ ์„ค์ •
import tempfile
from flask_session import Session
# ์ž„์‹œ ๋””๋ ‰ํ† ๋ฆฌ๋ฅผ ์‚ฌ์šฉํ•˜์—ฌ ๊ถŒํ•œ ๋ฌธ์ œ ํ•ด๊ฒฐ
session_dir = tempfile.gettempdir()
app.config['SESSION_TYPE'] = 'filesystem'
app.config['SESSION_PERMANENT'] = True
app.config['SESSION_USE_SIGNER'] = True
app.config['SESSION_FILE_DIR'] = session_dir
print(f"Using session directory: {session_dir}")
Session(app)
# ์‚ฌ์šฉ์ž ํด๋ž˜์Šค ์ •์˜
class User(UserMixin):
def __init__(self, id, username, password):
self.id = id
self.username = username
self.password = password
def get_id(self):
return str(self.id) # Flask-Login์€ ๋ฌธ์ž์—ด ID๋ฅผ ์š”๊ตฌํ•จ
# ํ…Œ์ŠคํŠธ์šฉ ์‚ฌ์šฉ์ž (์‹ค์ œ ํ™˜๊ฒฝ์—์„œ๋Š” ๋ฐ์ดํ„ฐ๋ฒ ์ด์Šค ์‚ฌ์šฉ ๊ถŒ์žฅ)
users = {
'admin': User('1', 'admin', 'admin123'),
'user': User('2', 'user', 'user123')
}
# ์‚ฌ์šฉ์ž ๋กœ๋” ํ•จ์ˆ˜
@login_manager.user_loader
def load_user(user_id):
print(f"Loading user with ID: {user_id}")
# ์„ธ์…˜ ๋””๋ฒ„๊ทธ ์ •๋ณด ์ถœ๋ ฅ
print(f"Session data in user_loader: {dict(session)}")
print(f"Current request cookies: {request.cookies}")
# user_id๊ฐ€ ๋ฌธ์ž์—ด๋กœ ์ „๋‹ฌ๋˜๊ธฐ ๋•Œ๋ฌธ์— ์‚ฌ์šฉ์ž ID๋กœ ์ฒ˜๋ฆฌ
for username, user in users.items():
if str(user.id) == str(user_id): # ํ™•์‹คํ•œ ๋ฌธ์ž์—ด ๋น„๊ต
print(f"User found: {username}, ID: {user.id}")
# ์„ธ์…˜ ์ •๋ณด ์—…๋ฐ์ดํŠธ
session['user_id'] = user.id
session['username'] = username
session.modified = True
return user
print(f"User not found with ID: {user_id}")
return None
# Model initialization
print("Loading models... This may take a moment.")
# Image embedding model (CLIP) for vector search
clip_model = None
clip_processor = None
try:
from transformers import CLIPProcessor, CLIPModel
# ์ž„์‹œ ๋””๋ ‰ํ† ๋ฆฌ ์‚ฌ์šฉ
import tempfile
temp_dir = tempfile.gettempdir()
os.environ["TRANSFORMERS_CACHE"] = temp_dir
# CLIP ๋ชจ๋ธ ๋กœ๋“œ (์ด๋ฏธ์ง€ ์ž„๋ฒ ๋”ฉ์šฉ)
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
print("CLIP model loaded successfully")
except Exception as e:
print("Error loading CLIP model:", e)
clip_model = None
clip_processor = None
# Vector DB ์ดˆ๊ธฐํ™”
vector_db = None
image_collection = None
object_collection = None
try:
# ChromaDB ํด๋ผ์ด์–ธํŠธ ์ดˆ๊ธฐํ™” (์ธ๋ฉ”๋ชจ๋ฆฌ DB)
vector_db = chromadb.Client()
# ์ž„๋ฒ ๋”ฉ ํ•จ์ˆ˜ ์„ค์ •
ef = embedding_functions.DefaultEmbeddingFunction()
# ์ด๋ฏธ์ง€ ์ปฌ๋ ‰์…˜ ์ƒ์„ฑ
image_collection = vector_db.create_collection(
name="image_collection",
embedding_function=ef,
get_or_create=True
)
# ๊ฐ์ฒด ์ธ์‹ ๊ฒฐ๊ณผ ์ปฌ๋ ‰์…˜ ์ƒ์„ฑ
object_collection = vector_db.create_collection(
name="object_collection",
embedding_function=ef,
get_or_create=True
)
print("Vector DB initialized successfully")
except Exception as e:
print("Error initializing Vector DB:", e)
vector_db = None
image_collection = None
object_collection = None
# YOLOv8 model
yolo_model = None
try:
import os
from ultralytics import YOLO
# ๋ชจ๋ธ ํŒŒ์ผ ๊ฒฝ๋กœ - ์ž„์‹œ ๋””๋ ‰ํ† ๋ฆฌ ์‚ฌ์šฉ
import tempfile
temp_dir = tempfile.gettempdir()
model_path = os.path.join(temp_dir, "yolov8n.pt")
# ๋ชจ๋ธ ํŒŒ์ผ์ด ์—†์œผ๋ฉด ์ง์ ‘ ๋‹ค์šด๋กœ๋“œ
if not os.path.exists(model_path):
print(f"Downloading YOLOv8 model to {model_path}...")
try:
os.system(f"wget -q https://ultralytics.com/assets/yolov8n.pt -O {model_path}")
print("YOLOv8 model downloaded successfully")
except Exception as e:
print(f"Error downloading YOLOv8 model: {e}")
# ๋‹ค์šด๋กœ๋“œ ์‹คํŒจ ์‹œ ๋Œ€์ฒด URL ์‹œ๋„
try:
os.system(f"wget -q https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt -O {model_path}")
print("YOLOv8 model downloaded from alternative source")
except Exception as e2:
print(f"Error downloading from alternative source: {e2}")
# ๋งˆ์ง€๋ง‰ ๋Œ€์•ˆ์œผ๋กœ ์ง์ ‘ ๋ชจ๋ธ URL ์‚ฌ์šฉ
try:
os.system(f"curl -L https://ultralytics.com/assets/yolov8n.pt --output {model_path}")
print("YOLOv8 model downloaded using curl")
except Exception as e3:
print(f"All download attempts failed: {e3}")
# ํ™˜๊ฒฝ ๋ณ€์ˆ˜ ์„ค์ • - ์„ค์ • ํŒŒ์ผ ๊ฒฝ๋กœ ์ง€์ •
os.environ["YOLO_CONFIG_DIR"] = temp_dir
os.environ["MPLCONFIGDIR"] = temp_dir
yolo_model = YOLO(model_path) # Using the nano model for faster inference
print("YOLOv8 model loaded successfully")
except Exception as e:
print("Error loading YOLOv8 model:", e)
yolo_model = None
# DETR model (DEtection TRansformer)
detr_processor = None
detr_model = None
try:
from transformers import DetrImageProcessor, DetrForObjectDetection
detr_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
detr_model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
print("DETR model loaded successfully")
except Exception as e:
print("Error loading DETR model:", e)
detr_processor = None
detr_model = None
# ViT model
vit_processor = None
vit_model = None
try:
from transformers import ViTImageProcessor, ViTForImageClassification
vit_processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
vit_model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
print("ViT model loaded successfully")
except Exception as e:
print("Error loading ViT model:", e)
vit_processor = None
vit_model = None
# Get device information
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# LLM model (using an open-access model instead of Llama 4 which requires authentication)
llm_model = None
llm_tokenizer = None
try:
from transformers import AutoModelForCausalLM, AutoTokenizer
print("Loading LLM model... This may take a moment.")
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # Using TinyLlama as an open-access alternative
llm_tokenizer = AutoTokenizer.from_pretrained(model_name)
llm_model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
# Removing options that require accelerate package
# device_map="auto",
# load_in_8bit=True
).to(device)
print("LLM model loaded successfully")
except Exception as e:
print(f"Error loading LLM model: {e}")
llm_model = None
llm_tokenizer = None
def process_llm_query(vision_results, user_query):
"""Process a query with the LLM model using vision results and user text"""
if llm_model is None or llm_tokenizer is None:
return {"error": "LLM model not available"}
# ๊ฒฐ๊ณผ ๋ฐ์ดํ„ฐ ์š”์•ฝ (ํ† ํฐ ๊ธธ์ด ์ œํ•œ์„ ์œ„ํ•ด)
summarized_results = []
# ๊ฐ์ฒด ํƒ์ง€ ๊ฒฐ๊ณผ ์š”์•ฝ
if isinstance(vision_results, list):
# ์ตœ๋Œ€ 10๊ฐœ ๊ฐ์ฒด๋งŒ ํฌํ•จ
for i, obj in enumerate(vision_results[:10]):
if isinstance(obj, dict):
# ํ•„์š”ํ•œ ์ •๋ณด๋งŒ ์ถ”์ถœ
summary = {
"label": obj.get("label", "unknown"),
"confidence": obj.get("confidence", 0),
}
summarized_results.append(summary)
# Create a prompt combining vision results and user query
prompt = f"""You are an AI assistant analyzing image detection results.
Here are the objects detected in the image: {json.dumps(summarized_results, indent=2)}
User question: {user_query}
Please provide a detailed analysis based on the detected objects and the user's question.
"""
# Tokenize and generate response
try:
start_time = time.time()
# ํ† ํฐ ๊ธธ์ด ํ™•์ธ ๋ฐ ์ œํ•œ
tokens = llm_tokenizer.encode(prompt)
if len(tokens) > 1500: # ์•ˆ์ „ ๋งˆ์ง„ ์„ค์ •
prompt = f"""You are an AI assistant analyzing image detection results.
The image contains {len(summarized_results)} detected objects.
User question: {user_query}
Please provide a general analysis based on the user's question.
"""
inputs = llm_tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
output = llm_model.generate(
**inputs,
max_new_tokens=512,
temperature=0.7,
top_p=0.9,
do_sample=True
)
response_text = llm_tokenizer.decode(output[0], skip_special_tokens=True)
# Remove the prompt from the response
if response_text.startswith(prompt):
response_text = response_text[len(prompt):].strip()
inference_time = time.time() - start_time
return {
"response": response_text,
"performance": {
"inference_time": round(inference_time, 3),
"device": "GPU" if torch.cuda.is_available() else "CPU"
}
}
except Exception as e:
return {"error": f"Error processing LLM query: {str(e)}"}
def image_to_base64(img):
"""Convert PIL Image to base64 string"""
buffered = io.BytesIO()
img.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
return img_str
def process_yolo(image):
if yolo_model is None:
return {"error": "YOLOv8 model not loaded"}
# Measure inference time
start_time = time.time()
# Convert to numpy if it's a PIL image
if isinstance(image, Image.Image):
image_np = np.array(image)
else:
image_np = image
# Run inference
results = yolo_model(image_np)
# Process results
result_image = results[0].plot()
result_image = Image.fromarray(result_image)
# Get detection information
boxes = results[0].boxes
class_names = results[0].names
# Format detection results
detections = []
for box in boxes:
class_id = int(box.cls[0].item())
class_name = class_names[class_id]
confidence = round(box.conf[0].item(), 2)
bbox = box.xyxy[0].tolist()
bbox = [round(x) for x in bbox]
detections.append({
"class": class_name,
"confidence": confidence,
"bbox": bbox
})
# Calculate inference time
inference_time = time.time() - start_time
# Add inference time and device info
device_info = "GPU" if torch.cuda.is_available() else "CPU"
return {
"image": image_to_base64(result_image),
"detections": detections,
"performance": {
"inference_time": round(inference_time, 3),
"device": device_info
}
}
def process_detr(image):
if detr_model is None or detr_processor is None:
return {"error": "DETR model not loaded"}
# Measure inference time
start_time = time.time()
# Prepare image for the model
inputs = detr_processor(images=image, return_tensors="pt")
# Run inference
with torch.no_grad():
outputs = detr_model(**inputs)
# Process results
target_sizes = torch.tensor([image.size[::-1]])
results = detr_processor.post_process_object_detection(
outputs, target_sizes=target_sizes, threshold=0.9
)[0]
# Create a copy of the image to draw on
result_image = image.copy()
fig, ax = plt.subplots(1)
ax.imshow(result_image)
# Format detection results
detections = []
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
box = [round(i) for i in box.tolist()]
class_name = detr_model.config.id2label[label.item()]
confidence = round(score.item(), 2)
# Draw rectangle
rect = Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1],
linewidth=2, edgecolor='r', facecolor='none')
ax.add_patch(rect)
# Add label
plt.text(box[0], box[1], "{}: {}".format(class_name, confidence),
bbox=dict(facecolor='white', alpha=0.8))
detections.append({
"class": class_name,
"confidence": confidence,
"bbox": box
})
# Save figure to image
buf = io.BytesIO()
plt.tight_layout()
plt.axis('off')
plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
buf.seek(0)
result_image = Image.open(buf)
plt.close(fig)
# Calculate inference time
inference_time = time.time() - start_time
# Add inference time and device info
device_info = "GPU" if torch.cuda.is_available() else "CPU"
return {
"image": image_to_base64(result_image),
"detections": detections,
"performance": {
"inference_time": round(inference_time, 3),
"device": device_info
}
}
def process_vit(image):
if vit_model is None or vit_processor is None:
return {"error": "ViT model not loaded"}
# Measure inference time
start_time = time.time()
# Prepare image for the model
inputs = vit_processor(images=image, return_tensors="pt")
# Run inference
with torch.no_grad():
outputs = vit_model(**inputs)
logits = outputs.logits
# Get the predicted class
predicted_class_idx = logits.argmax(-1).item()
prediction = vit_model.config.id2label[predicted_class_idx]
# Get top 5 predictions
probs = torch.nn.functional.softmax(logits, dim=-1)[0]
top5_prob, top5_indices = torch.topk(probs, 5)
results = []
for i, (prob, idx) in enumerate(zip(top5_prob, top5_indices)):
class_name = vit_model.config.id2label[idx.item()]
results.append({
"rank": i+1,
"class": class_name,
"probability": round(prob.item(), 3)
})
# Calculate inference time
inference_time = time.time() - start_time
# Add inference time and device info
device_info = "GPU" if torch.cuda.is_available() else "CPU"
return {
"top_predictions": results,
"performance": {
"inference_time": round(inference_time, 3),
"device": device_info
}
}
@app.route('/api/detect/yolo', methods=['POST'])
@login_required
def yolo_detect():
if 'image' not in request.files:
return jsonify({"error": "No image provided"}), 400
file = request.files['image']
image = Image.open(file.stream)
result = process_yolo(image)
return jsonify(result)
@app.route('/api/detect/detr', methods=['POST'])
@login_required
def detr_detect():
if 'image' not in request.files:
return jsonify({"error": "No image provided"}), 400
file = request.files['image']
image = Image.open(file.stream)
result = process_detr(image)
return jsonify(result)
@app.route('/api/classify/vit', methods=['POST'])
@login_required
def vit_classify():
if 'image' not in request.files:
return jsonify({"error": "No image provided"}), 400
file = request.files['image']
image = Image.open(file.stream)
result = process_vit(image)
return jsonify(result)
@app.route('/api/analyze', methods=['POST'])
@login_required
def analyze_with_llm():
# Check if required data is in the request
if not request.json:
return jsonify({"error": "No JSON data provided"}), 400
# Extract vision results and user query from request
data = request.json
if 'visionResults' not in data or 'userQuery' not in data:
return jsonify({"error": "Missing required fields: visionResults or userQuery"}), 400
vision_results = data['visionResults']
user_query = data['userQuery']
# Process the query with LLM
result = process_llm_query(vision_results, user_query)
return jsonify(result)
def generate_image_embedding(image):
"""CLIP ๋ชจ๋ธ์„ ์‚ฌ์šฉํ•˜์—ฌ ์ด๋ฏธ์ง€ ์ž„๋ฒ ๋”ฉ ์ƒ์„ฑ"""
if clip_model is None or clip_processor is None:
return None
try:
# ์ด๋ฏธ์ง€ ์ „์ฒ˜๋ฆฌ
inputs = clip_processor(images=image, return_tensors="pt")
# ์ด๋ฏธ์ง€ ์ž„๋ฒ ๋”ฉ ์ƒ์„ฑ
with torch.no_grad():
image_features = clip_model.get_image_features(**inputs)
# ์ž„๋ฒ ๋”ฉ ์ •๊ทœํ™” ๋ฐ numpy ๋ฐฐ์—ด๋กœ ๋ณ€ํ™˜
image_embedding = image_features.squeeze().cpu().numpy()
normalized_embedding = image_embedding / np.linalg.norm(image_embedding)
return normalized_embedding.tolist()
except Exception as e:
print(f"Error generating image embedding: {e}")
return None
@app.route('/api/similar-images', methods=['POST'])
@login_required
def find_similar_images():
"""์œ ์‚ฌ ์ด๋ฏธ์ง€ ๊ฒ€์ƒ‰ API"""
if clip_model is None or clip_processor is None or image_collection is None:
return jsonify({"error": "Image embedding model or vector DB not available"})
try:
# ์š”์ฒญ์—์„œ ์ด๋ฏธ์ง€ ๋ฐ์ดํ„ฐ ์ถ”์ถœ
if 'image' not in request.files and 'image' not in request.form:
return jsonify({"error": "No image provided"})
if 'image' in request.files:
# ํŒŒ์ผ๋กœ ์—…๋กœ๋“œ๋œ ๊ฒฝ์šฐ
image_file = request.files['image']
image = Image.open(image_file).convert('RGB')
else:
# base64๋กœ ์ธ์ฝ”๋”ฉ๋œ ๊ฒฝ์šฐ
image_data = request.form['image']
if image_data.startswith('data:image'):
# Remove the data URL prefix if present
image_data = image_data.split(',')[1]
image = Image.open(BytesIO(base64.b64decode(image_data))).convert('RGB')
# ์ด๋ฏธ์ง€ ID ์ƒ์„ฑ (์ž„์‹œ)
image_id = str(uuid.uuid4())
# ์ด๋ฏธ์ง€ ์ž„๋ฒ ๋”ฉ ์ƒ์„ฑ
embedding = generate_image_embedding(image)
if embedding is None:
return jsonify({"error": "Failed to generate image embedding"})
# ํ˜„์žฌ ์ด๋ฏธ์ง€๋ฅผ DB์— ์ถ”๊ฐ€ (์„ ํƒ์ )
# image_collection.add(
# ids=[image_id],
# embeddings=[embedding]
# )
# ์œ ์‚ฌ ์ด๋ฏธ์ง€ ๊ฒ€์ƒ‰
results = image_collection.query(
query_embeddings=[embedding],
n_results=5 # ์ƒ์œ„ 5๊ฐœ ๊ฒฐ๊ณผ ๋ฐ˜ํ™˜
)
# ๊ฒฐ๊ณผ ํฌ๋งทํŒ…
similar_images = []
if len(results['ids'][0]) > 0:
for i, img_id in enumerate(results['ids'][0]):
similar_images.append({
"id": img_id,
"distance": float(results['distances'][0][i]) if 'distances' in results else 0.0,
"metadata": results['metadatas'][0][i] if 'metadatas' in results else {}
})
return jsonify({
"query_image_id": image_id,
"similar_images": similar_images
})
except Exception as e:
print(f"Error in similar-images API: {e}")
return jsonify({"error": str(e)}), 500
@app.route('/api/add-to-collection', methods=['POST'])
@login_required
def add_to_collection():
"""์ด๋ฏธ์ง€๋ฅผ ๋ฒกํ„ฐ DB์— ์ถ”๊ฐ€ํ•˜๋Š” API"""
if clip_model is None or clip_processor is None or image_collection is None:
return jsonify({"error": "Image embedding model or vector DB not available"})
try:
# ์š”์ฒญ์—์„œ ์ด๋ฏธ์ง€ ๋ฐ์ดํ„ฐ ์ถ”์ถœ
if 'image' not in request.files and 'image' not in request.form:
return jsonify({"error": "No image provided"})
# ๋ฉ”ํƒ€๋ฐ์ดํ„ฐ ์ถ”์ถœ
metadata = {}
if 'metadata' in request.form:
metadata = json.loads(request.form['metadata'])
# ์ด๋ฏธ์ง€ ID (์ œ๊ณต๋˜์ง€ ์•Š์€ ๊ฒฝ์šฐ ์ž๋™ ์ƒ์„ฑ)
image_id = request.form.get('id', str(uuid.uuid4()))
if 'image' in request.files:
# ํŒŒ์ผ๋กœ ์—…๋กœ๋“œ๋œ ๊ฒฝ์šฐ
image_file = request.files['image']
image = Image.open(image_file).convert('RGB')
else:
# base64๋กœ ์ธ์ฝ”๋”ฉ๋œ ๊ฒฝ์šฐ
image_data = request.form['image']
if image_data.startswith('data:image'):
# Remove the data URL prefix if present
image_data = image_data.split(',')[1]
image = Image.open(BytesIO(base64.b64decode(image_data))).convert('RGB')
# ์ด๋ฏธ์ง€ ์ž„๋ฒ ๋”ฉ ์ƒ์„ฑ
embedding = generate_image_embedding(image)
if embedding is None:
return jsonify({"error": "Failed to generate image embedding"})
# ์ด๋ฏธ์ง€ ๋ฐ์ดํ„ฐ๋ฅผ base64๋กœ ์ธ์ฝ”๋”ฉํ•˜์—ฌ ๋ฉ”ํƒ€๋ฐ์ดํ„ฐ์— ์ถ”๊ฐ€
buffered = BytesIO()
image.save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
metadata['image_data'] = img_str
# ์ด๋ฏธ์ง€๋ฅผ DB์— ์ถ”๊ฐ€
image_collection.add(
ids=[image_id],
embeddings=[embedding],
metadatas=[metadata]
)
return jsonify({
"success": True,
"image_id": image_id,
"message": "Image added to collection"
})
except Exception as e:
print(f"Error in add-to-collection API: {e}")
return jsonify({"error": str(e)}), 500
@app.route('/api/add-detected-objects', methods=['POST'])
@login_required
def add_detected_objects():
"""๊ฐ์ฒด ์ธ์‹ ๊ฒฐ๊ณผ๋ฅผ ๋ฒกํ„ฐ DB์— ์ถ”๊ฐ€ํ•˜๋Š” API"""
if clip_model is None or object_collection is None:
return jsonify({"error": "Image embedding model or vector DB not available"})
try:
# ๋””๋ฒ„๊น…: ์š”์ฒญ ๋ฐ์ดํ„ฐ ๋กœ๊น…
print("[DEBUG] Received request in add-detected-objects")
# ์š”์ฒญ์—์„œ ์ด๋ฏธ์ง€์™€ ๊ฐ์ฒด ๊ฒ€์ถœ ๊ฒฐ๊ณผ ๋ฐ์ดํ„ฐ ์ถ”์ถœ
data = request.json
print(f"[DEBUG] Request data keys: {list(data.keys()) if data else 'None'}")
if not data:
print("[DEBUG] Error: No data received in request")
return jsonify({"error": "No data received"})
if 'image' not in data:
print("[DEBUG] Error: 'image' key not found in request data")
return jsonify({"error": "Missing image data"})
if 'objects' not in data:
print("[DEBUG] Error: 'objects' key not found in request data")
return jsonify({"error": "Missing objects data"})
# ์ด๋ฏธ์ง€ ๋ฐ์ดํ„ฐ ๋””๋ฒ„๊น…
print(f"[DEBUG] Image data type: {type(data['image'])}")
print(f"[DEBUG] Image data starts with: {data['image'][:50]}...") # ์ฒ˜์Œ 50์ž๋งŒ ์ถœ๋ ฅ
# ๊ฐ์ฒด ๋ฐ์ดํ„ฐ ๋””๋ฒ„๊น…
print(f"[DEBUG] Objects data type: {type(data['objects'])}")
print(f"[DEBUG] Objects count: {len(data['objects']) if isinstance(data['objects'], list) else 'Not a list'}")
if isinstance(data['objects'], list) and len(data['objects']) > 0:
print(f"[DEBUG] First object keys: {list(data['objects'][0].keys()) if isinstance(data['objects'][0], dict) else 'Not a dict'}")
# ์ด๋ฏธ์ง€ ๋ฐ์ดํ„ฐ ์ฒ˜๋ฆฌ
image_data = data['image']
if image_data.startswith('data:image'):
image_data = image_data.split(',')[1]
image = Image.open(BytesIO(base64.b64decode(image_data))).convert('RGB')
image_width, image_height = image.size
# ์ด๋ฏธ์ง€ ID
image_id = data.get('imageId', str(uuid.uuid4()))
# ๊ฐ์ฒด ๋ฐ์ดํ„ฐ ์ฒ˜๋ฆฌ
objects = data['objects']
object_ids = []
object_embeddings = []
object_metadatas = []
for obj in objects:
# ๊ฐ์ฒด ID ์ƒ์„ฑ
object_id = f"{image_id}_{str(uuid.uuid4())[:8]}"
# ๋ฐ”์šด๋”ฉ ๋ฐ•์Šค ์ •๋ณด ์ถ”์ถœ
bbox = obj.get('bbox', [])
# ๋ฆฌ์ŠคํŠธ ํ˜•ํƒœ์˜ bbox [x1, y1, x2, y2] ์ฒ˜๋ฆฌ
if isinstance(bbox, list) and len(bbox) >= 4:
x1 = bbox[0] / image_width # ์ •๊ทœํ™”๋œ ์ขŒํ‘œ๋กœ ๋ณ€ํ™˜
y1 = bbox[1] / image_height
x2 = bbox[2] / image_width
y2 = bbox[3] / image_height
width = x2 - x1
height = y2 - y1
# ๋”•์…”๋„ˆ๋ฆฌ ํ˜•ํƒœ์˜ bbox {'x': x, 'y': y, 'width': width, 'height': height} ์ฒ˜๋ฆฌ
elif isinstance(bbox, dict):
x1 = bbox.get('x', 0)
y1 = bbox.get('y', 0)
width = bbox.get('width', 0)
height = bbox.get('height', 0)
else:
# ๊ธฐ๋ณธ๊ฐ’ ์„ค์ •
x1, y1, width, height = 0, 0, 0, 0
# ๋ฐ”์šด๋”ฉ ๋ฐ•์Šค๋ฅผ ์ด๋ฏธ์ง€ ์ขŒํ‘œ๋กœ ๋ณ€ํ™˜
x1_px = int(x1 * image_width)
y1_px = int(y1 * image_height)
width_px = int(width * image_width)
height_px = int(height * image_height)
# ๊ฐ์ฒด ์ด๋ฏธ์ง€ ์ž๋ฅด๊ธฐ
try:
object_image = image.crop((x1_px, y1_px, x1_px + width_px, y1_px + height_px))
# ์ž„๋ฒ ๋”ฉ ์ƒ์„ฑ
embedding = generate_image_embedding(object_image)
if embedding is None:
continue
# ๋ฉ”ํƒ€๋ฐ์ดํ„ฐ ๊ตฌ์„ฑ
# bbox๋ฅผ JSON ๋ฌธ์ž์—ด๋กœ ์ง๋ ฌํ™”ํ•˜์—ฌ ChromaDB ๋ฉ”ํƒ€๋ฐ์ดํ„ฐ ์ œํ•œ ์šฐํšŒ
bbox_json = json.dumps({
"x": x1,
"y": y1,
"width": width,
"height": height
})
# ๊ฐ์ฒด ์ด๋ฏธ์ง€๋ฅผ base64๋กœ ์ธ์ฝ”๋”ฉ
buffered = BytesIO()
object_image.save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
metadata = {
"image_id": image_id,
"class": obj.get('class', ''),
"confidence": obj.get('confidence', 0),
"bbox": bbox_json, # JSON ๋ฌธ์ž์—ด๋กœ ์ €์žฅ
"image_data": img_str # ์ด๋ฏธ์ง€ ๋ฐ์ดํ„ฐ ์ถ”๊ฐ€
}
object_ids.append(object_id)
object_embeddings.append(embedding)
object_metadatas.append(metadata)
except Exception as e:
print(f"Error processing object: {e}")
continue
# ๊ฐ์ฒด๊ฐ€ ์—†๋Š” ๊ฒฝ์šฐ
if not object_ids:
return jsonify({"error": "No valid objects to add"})
# ๋””๋ฒ„๊น…: ๋ฉ”ํƒ€๋ฐ์ดํ„ฐ ์ถœ๋ ฅ
print(f"[DEBUG] Adding {len(object_ids)} objects to vector DB")
print(f"[DEBUG] First metadata sample: {object_metadatas[0] if object_metadatas else 'None'}")
try:
# ๊ฐ์ฒด๋“ค์„ DB์— ์ถ”๊ฐ€
object_collection.add(
ids=object_ids,
embeddings=object_embeddings,
metadatas=object_metadatas
)
print("[DEBUG] Successfully added objects to vector DB")
except Exception as e:
print(f"[DEBUG] Error adding to vector DB: {e}")
raise e
return jsonify({
"success": True,
"image_id": image_id,
"object_count": len(object_ids),
"object_ids": object_ids
})
except Exception as e:
print(f"Error in add-detected-objects API: {e}")
return jsonify({"error": str(e)}), 500
@app.route('/api/search-similar-objects', methods=['POST'])
@login_required
def search_similar_objects():
"""์œ ์‚ฌํ•œ ๊ฐ์ฒด ๊ฒ€์ƒ‰ API"""
print("[DEBUG] Received request in search-similar-objects")
if clip_model is None or object_collection is None:
print("[DEBUG] Error: Image embedding model or vector DB not available")
return jsonify({"error": "Image embedding model or vector DB not available"})
try:
# ์š”์ฒญ ๋ฐ์ดํ„ฐ ์ถ”์ถœ
data = request.json
print(f"[DEBUG] Request data keys: {list(data.keys()) if data else 'None'}")
if not data:
print("[DEBUG] Error: Missing request data")
return jsonify({"error": "Missing request data"})
# ๊ฒ€์ƒ‰ ์œ ํ˜• ๊ฒฐ์ •
search_type = data.get('searchType', 'image')
n_results = int(data.get('n_results', 5)) # ๊ฒฐ๊ณผ ๊ฐœ์ˆ˜
print(f"[DEBUG] Search type: {search_type}, n_results: {n_results}")
query_embedding = None
if search_type == 'image' and 'image' in data:
# ์ด๋ฏธ์ง€๋กœ ๊ฒ€์ƒ‰ํ•˜๋Š” ๊ฒฝ์šฐ
print("[DEBUG] Searching by image")
image_data = data['image']
if image_data.startswith('data:image'):
image_data = image_data.split(',')[1]
try:
image = Image.open(BytesIO(base64.b64decode(image_data))).convert('RGB')
query_embedding = generate_image_embedding(image)
print(f"[DEBUG] Generated image embedding: {type(query_embedding)}, shape: {len(query_embedding) if query_embedding is not None else 'None'}")
except Exception as e:
print(f"[DEBUG] Error generating image embedding: {e}")
return jsonify({"error": f"Error processing image: {str(e)}"}), 500
elif search_type == 'object' and 'objectId' in data:
# ๊ฐ์ฒด ID๋กœ ๊ฒ€์ƒ‰ํ•˜๋Š” ๊ฒฝ์šฐ
object_id = data['objectId']
result = object_collection.get(ids=[object_id], include=["embeddings"])
if result and "embeddings" in result and len(result["embeddings"]) > 0:
query_embedding = result["embeddings"][0]
elif search_type == 'class' and 'class_name' in data:
# ํด๋ž˜์Šค ์ด๋ฆ„์œผ๋กœ ๊ฒ€์ƒ‰ํ•˜๋Š” ๊ฒฝ์šฐ
print("[DEBUG] Searching by class name")
class_name = data['class_name']
print(f"[DEBUG] Class name: {class_name}")
filter_query = {"class": {"$eq": class_name}}
try:
# ํด๋ž˜์Šค๋กœ ํ•„ํ„ฐ๋งํ•˜์—ฌ ๊ฒ€์ƒ‰
print(f"[DEBUG] Querying with filter: {filter_query}")
# Use get method instead of query for class-based filtering
results = object_collection.get(
where=filter_query,
limit=n_results,
include=["metadatas", "embeddings", "documents"]
)
print(f"[DEBUG] Query results: {results['ids'][0] if 'ids' in results and len(results['ids']) > 0 else 'No results'}")
formatted_results = format_object_results(results)
print(f"[DEBUG] Formatted results count: {len(formatted_results)}")
return jsonify({
"success": True,
"searchType": "class",
"results": formatted_results
})
except Exception as e:
print(f"[DEBUG] Error in class search: {e}")
return jsonify({"error": f"Error in class search: {str(e)}"}), 500
else:
print(f"[DEBUG] Invalid search parameters: {data}")
return jsonify({"error": "Invalid search parameters"})
if query_embedding is None:
print("[DEBUG] Error: Failed to generate query embedding")
return jsonify({"error": "Failed to generate query embedding"})
try:
# ์œ ์‚ฌ๋„ ๊ฒ€์ƒ‰ ์‹คํ–‰
print(f"[DEBUG] Running similarity search with embedding of length {len(query_embedding)}")
results = object_collection.query(
query_embeddings=[query_embedding],
n_results=n_results,
include=["metadatas", "distances"]
)
print(f"[DEBUG] Query results: {results['ids'][0] if 'ids' in results and len(results['ids']) > 0 else 'No results'}")
formatted_results = format_object_results(results)
print(f"[DEBUG] Formatted results count: {len(formatted_results)}")
return jsonify({
"success": True,
"searchType": search_type,
"results": formatted_results
})
except Exception as e:
print(f"[DEBUG] Error in similarity search: {e}")
return jsonify({"error": f"Error in similarity search: {str(e)}"}), 500
except Exception as e:
print(f"Error in search-similar-objects API: {e}")
return jsonify({"error": str(e)}), 500
def format_object_results(results):
"""๊ฒ€์ƒ‰ ๊ฒฐ๊ณผ ํฌ๋งทํŒ… - ChromaDB query ๋ฐ get ๋ฉ”์„œ๋“œ ๊ฒฐ๊ณผ ๋ชจ๋‘ ์ฒ˜๋ฆฌ"""
formatted_results = []
print(f"[DEBUG] Formatting results: {results.keys() if results else 'None'}")
if not results:
print("[DEBUG] No results to format")
return formatted_results
try:
# Check if this is a result from 'get' method (class search) or 'query' method (similarity search)
is_get_result = 'ids' in results and isinstance(results['ids'], list) and not isinstance(results['ids'][0], list) if 'ids' in results else False
if is_get_result:
# Handle results from 'get' method (flat structure)
print("[DEBUG] Processing results from get method (class search)")
if len(results['ids']) == 0:
return formatted_results
for i, obj_id in enumerate(results['ids']):
try:
# Extract object info
metadata = results['metadatas'][i] if 'metadatas' in results else {}
# Deserialize bbox if stored as JSON string
if 'bbox' in metadata and isinstance(metadata['bbox'], str):
try:
metadata['bbox'] = json.loads(metadata['bbox'])
except:
pass # Keep as is if deserialization fails
result_item = {
"id": obj_id,
"metadata": metadata
}
# No distance in get results
# Check if image data is already in metadata
if 'image_data' not in metadata:
print(f"[DEBUG] Image data not found in metadata for object {obj_id}")
else:
print(f"[DEBUG] Image data found in metadata for object {obj_id}")
formatted_results.append(result_item)
except Exception as e:
print(f"[DEBUG] Error formatting get result {i}: {e}")
else:
# Handle results from 'query' method (nested structure)
print("[DEBUG] Processing results from query method (similarity search)")
if 'ids' not in results or len(results['ids']) == 0 or len(results['ids'][0]) == 0:
return formatted_results
for i, obj_id in enumerate(results['ids'][0]):
try:
# Extract object info
metadata = results['metadatas'][0][i] if 'metadatas' in results and len(results['metadatas']) > 0 else {}
# Deserialize bbox if stored as JSON string
if 'bbox' in metadata and isinstance(metadata['bbox'], str):
try:
metadata['bbox'] = json.loads(metadata['bbox'])
except:
pass # Keep as is if deserialization fails
result_item = {
"id": obj_id,
"metadata": metadata
}
if 'distances' in results and len(results['distances']) > 0:
result_item["distance"] = float(results['distances'][0][i])
# Check if image data is already in metadata
if 'image_data' not in metadata:
try:
# Try to get original image via image ID
image_id = metadata.get('image_id')
if image_id:
print(f"[DEBUG] Image data not found in metadata for object {obj_id} with image_id {image_id}")
except Exception as e:
print(f"[DEBUG] Error checking image data for result {i}: {e}")
else:
print(f"[DEBUG] Image data found in metadata for object {obj_id}")
formatted_results.append(result_item)
except Exception as e:
print(f"[DEBUG] Error formatting query result {i}: {e}")
except Exception as e:
print(f"[DEBUG] Error in format_object_results: {e}")
return formatted_results
# ๋กœ๊ทธ์ธ ํŽ˜์ด์ง€ HTML ํ…œํ”Œ๋ฆฟ
LOGIN_TEMPLATE = '''
<!DOCTYPE html>
<html lang="ko">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Vision LLM Agent - ๋กœ๊ทธ์ธ</title>
<style>
body {
font-family: Arial, sans-serif;
background-color: #f5f5f5;
display: flex;
justify-content: center;
align-items: center;
height: 100vh;
margin: 0;
}
.login-container {
background-color: white;
padding: 2rem;
border-radius: 8px;
box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
width: 100%;
max-width: 400px;
}
h1 {
text-align: center;
color: #4a6cf7;
margin-bottom: 1.5rem;
}
.form-group {
margin-bottom: 1rem;
}
label {
display: block;
margin-bottom: 0.5rem;
font-weight: bold;
}
input {
width: 100%;
padding: 0.75rem;
border: 1px solid #ddd;
border-radius: 4px;
font-size: 1rem;
}
button {
width: 100%;
padding: 0.75rem;
background-color: #4a6cf7;
color: white;
border: none;
border-radius: 4px;
font-size: 1rem;
cursor: pointer;
margin-top: 1rem;
}
button:hover {
background-color: #3a5cd8;
}
.error-message {
color: #e74c3c;
margin-top: 1rem;
text-align: center;
}
</style>
</head>
<body>
<div class="login-container">
<h1>Vision LLM Agent</h1>
<form action="/login" method="post" autocomplete="off">
<!-- hidden dummy fields to discourage Chrome autofill -->
<input type="text" name="fakeusernameremembered" style="display:none" tabindex="-1" autocomplete="off">
<input type="password" name="fakepasswordremembered" style="display:none" tabindex="-1" autocomplete="off">
<div class="form-group">
<label for="username">Username</label>
<input type="text" id="username" name="username" required autocomplete="username" autocapitalize="none" autocorrect="off" spellcheck="false">
</div>
<div class="form-group">
<label for="password">Password</label>
<input type="password" id="password" name="password" required autocomplete="current-password" autocapitalize="none" autocorrect="off" spellcheck="false">
</div>
<button type="submit">Login</button>
{% if error %}
<p class="error-message">{{ error }}</p>
{% endif %}
</form>
</div>
</body>
</html>
'''
@app.route('/login', methods=['GET', 'POST'])
def login():
# ์ด๋ฏธ ๋กœ๊ทธ์ธ๋œ ์‚ฌ์šฉ์ž๋Š” ๋ฉ”์ธ ํŽ˜์ด์ง€๋กœ ๋ฆฌ๋””๋ ‰์…˜
if current_user.is_authenticated and login_fresh():
print(f"User already authenticated and fresh as: {current_user.username}, redirecting to index")
return redirect('/index.html')
elif current_user.is_authenticated and not login_fresh():
# Remember-cookie ์ƒํƒœ ๋“ฑ ๋น„-ํ”„๋ ˆ์‹œ ์„ธ์…˜์ด๋ฉด ๋กœ๊ทธ์ธ ํŽ˜์ด์ง€๋ฅผ ๋ณด์—ฌ์„œ ์žฌ์ธ์ฆ ์œ ๋„
print("User authenticated but session not fresh; showing login page for reauthentication")
error = None
if request.method == 'POST':
username = request.form.get('username')
password = request.form.get('password')
print(f"Login attempt: username={username}")
if username in users and users[username].password == password:
# ๋กœ๊ทธ์ธ ์„ฑ๊ณต ์‹œ ์„ธ์…˜์— ์‚ฌ์šฉ์ž ์ •๋ณด ์ €์žฅ
user = users[username]
login_user(user, remember=False) # 2๋ถ„ ์„ธ์…˜ ๋งŒ๋ฃŒ๋ฅผ ์œ„ํ•ด remember ๋น„ํ™œ์„ฑํ™”
session['user_id'] = user.id
session['username'] = username
session.permanent = True
session.modified = True # ์„ธ์…˜ ๋ณ€๊ฒฝ ์‚ฌํ•ญ ์ฆ‰์‹œ ์ ์šฉ
print(f"Login successful for user: {username}, ID: {user.id}")
# ๋ฆฌ๋””๋ ‰์…˜ ์ฒ˜๋ฆฌ
next_page = request.args.get('next')
if next_page and next_page.startswith('/') and next_page != '/login':
print(f"Redirecting to: {next_page}")
return redirect(next_page)
print("Redirecting to index.html")
return redirect(url_for('serve_index_html'))
else:
error = 'Invalid username or password'
print(f"Login failed: {error}")
return render_template_string(LOGIN_TEMPLATE, error=error)
@app.route('/logout')
def logout():
logout_user()
# Clear server-side session fully
try:
session.clear()
except Exception as e:
print(f"[DEBUG] Error clearing session on logout: {e}")
# Ensure remember cookie is removed by setting an expired cookie
resp = redirect(url_for('login'))
try:
resp.delete_cookie(
key='remember_token',
path='/',
samesite='None',
secure=True,
httponly=True,
)
except Exception as e:
print(f"[DEBUG] Error deleting remember_token cookie: {e}")
return resp
# ์ •์  ํŒŒ์ผ ์„œ๋น™์„ ์œ„ํ•œ ๋ผ์šฐํŠธ (๋กœ๊ทธ์ธ ๋ถˆํ•„์š”)
@app.route('/static/<path:filename>')
def serve_static(filename):
print(f"Serving static file: {filename}")
resp = send_from_directory(app.static_folder, filename)
# Prevent caching of static assets to reflect latest frontend changes
resp.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
resp.headers['Pragma'] = 'no-cache'
resp.headers['Expires'] = '0'
return resp
# ์ธ๋ฑ์Šค HTML ์ง์ ‘ ์„œ๋น™ (๋กœ๊ทธ์ธ ํ•„์š”)
@app.route('/index.html')
def serve_index_html():
# ์„ธ์…˜ ๋ฐ ์ฟ ํ‚ค ๋””๋ฒ„๊ทธ ์ •๋ณด
print(f"Request to /index.html - Session data: {dict(session)}")
print(f"Request to /index.html - Cookies: {request.cookies}")
print(f"Request to /index.html - User authenticated: {current_user.is_authenticated}")
# ์ธ์ฆ ํ™•์ธ (fresh session only)
if not current_user.is_authenticated or not login_fresh():
print("User not authenticated, redirecting to login")
return redirect(url_for('login'))
print(f"Serving index.html for authenticated user: {current_user.username} (ID: {current_user.id})")
# ์„ธ์…˜ ์ƒํƒœ ๋””๋ฒ„๊ทธ
print(f"Session data: user_id={session.get('user_id')}, username={session.get('username')}, is_permanent={session.get('permanent', False)}")
# ์„ธ์…˜ ๋งŒ๋ฃŒ๋ฅผ ์˜๋„๋Œ€๋กœ ์œ ์ง€ํ•˜๊ธฐ ์œ„ํ•ด ์—ฌ๊ธฐ์„œ ์„ธ์…˜์„ ๊ฐฑ์‹ ํ•˜์ง€ ์•Š์Šต๋‹ˆ๋‹ค.
# ์ฃผ์˜: ์„ธ์…˜์— ์“ฐ๊ธฐ(๋˜๋Š” session.modified=True)๋Š” Flask-Session์—์„œ ๋งŒ๋ฃŒ์‹œ๊ฐ„์„ ์—ฐ์žฅํ•  ์ˆ˜ ์žˆ์Šต๋‹ˆ๋‹ค.
# index.html์„ ์ฝ์–ด ํ•˜ํŠธ๋น„ํŠธ ์Šคํฌ๋ฆฝํŠธ๋ฅผ ์ฃผ์ž…
index_path = os.path.join(app.static_folder, 'index.html')
try:
with open(index_path, 'r', encoding='utf-8') as f:
html = f.read()
except Exception as e:
print(f"[DEBUG] Failed to read index.html for injection: {e}")
return send_from_directory(app.static_folder, 'index.html')
heartbeat_script = """
<script>
(function(){
// 1) ์„ธ์…˜ ์ƒํƒœ ์ฃผ๊ธฐ ์ฒดํฌ (๋งŒ๋ฃŒ์‹œ ๋กœ๊ทธ์ธ์œผ๋กœ)
function checkSession(){
fetch('/api/status', {credentials: 'include', redirect: 'manual'}).then(function(res){
var redirected = res.redirected || (res.url && res.url.indexOf('/login') !== -1);
if(res.status !== 200 || redirected){
window.location.href = '/login';
}
}).catch(function(){
// ๋„คํŠธ์›Œํฌ ์˜ค๋ฅ˜ ๋“ฑ๋„ ๋กœ๊ทธ์ธ์œผ๋กœ ์œ ๋„
window.location.href = '/login';
});
}
checkSession();
setInterval(checkSession, 30000);
// 2) ์‚ฌ์šฉ์ž ๋น„ํ™œ์„ฑ(๋ฌด๋™์ž‘) 2๋ถ„ ํ›„ ์ž๋™ ๋กœ๊ทธ์•„์›ƒ
var idleMs = 120000; // 2๋ถ„
var idleTimer;
function triggerLogout(){
// ์„œ๋ฒ„ ์„ธ์…˜ ์ •๋ฆฌ ํ›„ ๋กœ๊ทธ์ธ ํ™”๋ฉด์œผ๋กœ
window.location.href = '/logout';
}
function resetIdle(){
if (idleTimer) clearTimeout(idleTimer);
idleTimer = setTimeout(triggerLogout, idleMs);
}
['click','mousemove','keydown','scroll','touchstart','visibilitychange'].forEach(function(evt){
window.addEventListener(evt, resetIdle, {passive:true});
});
resetIdle();
})();
</script>
"""
try:
if '</body>' in html:
html = html.replace('</body>', heartbeat_script + '</body>')
else:
html = html + heartbeat_script
except Exception as e:
print(f"[DEBUG] Failed to inject heartbeat script: {e}")
return send_from_directory(app.static_folder, 'index.html')
resp = make_response(html)
# Prevent sensitive pages from being cached
resp.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
resp.headers['Pragma'] = 'no-cache'
resp.headers['Expires'] = '0'
return resp
# Static files should be accessible without login requirements
@app.route('/static/<path:filename>')
def static_files(filename):
print(f"Serving static file: {filename}")
# Two possible locations after CRA build copy:
# 1) Top-level: static/<filename>
# 2) Nested build: static/static/<filename>
top_level_path = os.path.join(app.static_folder, filename)
nested_dir = os.path.join(app.static_folder, 'static')
nested_path = os.path.join(nested_dir, filename)
try:
if os.path.exists(top_level_path):
return send_from_directory(app.static_folder, filename)
elif os.path.exists(nested_path):
# Serve from nested build directory
return send_from_directory(nested_dir, filename)
else:
# Fallback: try as-is (may help in some edge cases)
return send_from_directory(app.static_folder, filename)
except Exception as e:
print(f"[DEBUG] Error serving static file '{filename}': {e}")
# Final fallback to avoid leaking stack traces
return ('Not Found', 404)
# Add explicit handlers for JS files that are failing
@app.route('/static/js/<path:filename>')
def static_js_files(filename):
print(f"Serving JS file: {filename}")
# Try top-level static/js and nested static/static/js
top_js_dir = os.path.join(app.static_folder, 'js')
nested_js_dir = os.path.join(app.static_folder, 'static', 'js')
top_js_path = os.path.join(top_js_dir, filename)
nested_js_path = os.path.join(nested_js_dir, filename)
try:
if os.path.exists(top_js_path):
return send_from_directory(top_js_dir, filename)
elif os.path.exists(nested_js_path):
return send_from_directory(nested_js_dir, filename)
else:
# As a fallback, let the generic static handler try
return static_files(os.path.join('js', filename))
except Exception as e:
print(f"[DEBUG] Error serving JS file '{filename}': {e}")
return ('Not Found', 404)
# ๊ธฐ๋ณธ ๊ฒฝ๋กœ ๋ฐ ๊ธฐํƒ€ ๊ฒฝ๋กœ ์ฒ˜๋ฆฌ (๋กœ๊ทธ์ธ ํ•„์š”)
@app.route('/', defaults={'path': ''}, methods=['GET'])
@app.route('/<path:path>', methods=['GET'])
@fresh_login_required
def serve_react(path):
"""Serve React frontend"""
print(f"Serving React frontend for path: {path}, user: {current_user.username if current_user.is_authenticated else 'not authenticated'}")
# ์ •์  ํŒŒ์ผ ์ฒ˜๋ฆฌ๋Š” ์ด์ œ ๋ณ„๋„ ๋ผ์šฐํŠธ์—์„œ ์ฒ˜๋ฆฌ
if path != "" and os.path.exists(os.path.join(app.static_folder, path)):
resp = send_from_directory(app.static_folder, path)
resp.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
resp.headers['Pragma'] = 'no-cache'
resp.headers['Expires'] = '0'
return resp
else:
# React ์•ฑ์˜ index.html ์„œ๋น™ (ํ•˜ํŠธ๋น„ํŠธ ์Šคํฌ๋ฆฝํŠธ ์ฃผ์ž…)
index_path = os.path.join(app.static_folder, 'index.html')
try:
with open(index_path, 'r', encoding='utf-8') as f:
html = f.read()
except Exception as e:
print(f"[DEBUG] Failed to read index.html for injection (serve_react): {e}")
resp = send_from_directory(app.static_folder, 'index.html')
resp.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
resp.headers['Pragma'] = 'no-cache'
resp.headers['Expires'] = '0'
return resp
heartbeat_script = """
<script>
(function(){
// 1) ์„ธ์…˜ ์ƒํƒœ ์ฃผ๊ธฐ ์ฒดํฌ (๋งŒ๋ฃŒ์‹œ ๋กœ๊ทธ์ธ์œผ๋กœ)
function checkSession(){
fetch('/api/status', {credentials: 'include', redirect: 'manual'}).then(function(res){
var redirected = res.redirected || (res.url && res.url.indexOf('/login') !== -1);
if(res.status !== 200 || redirected){
window.location.href = '/login';
}
}).catch(function(){
window.location.href = '/login';
});
}
checkSession();
setInterval(checkSession, 30000);
// 2) ์‚ฌ์šฉ์ž ๋น„ํ™œ์„ฑ(๋ฌด๋™์ž‘) 2๋ถ„ ํ›„ ์ž๋™ ๋กœ๊ทธ์•„์›ƒ
var idleMs = 120000; // 2๋ถ„
var idleTimer;
function triggerLogout(){
window.location.href = '/logout';
}
function resetIdle(){
if (idleTimer) clearTimeout(idleTimer);
idleTimer = setTimeout(triggerLogout, idleMs);
}
['click','mousemove','keydown','scroll','touchstart','visibilitychange'].forEach(function(evt){
window.addEventListener(evt, resetIdle, {passive:true});
});
resetIdle();
})();
</script>
"""
try:
if '</body>' in html:
html = html.replace('</body>', heartbeat_script + '</body>')
else:
html = html + heartbeat_script
except Exception as e:
print(f"[DEBUG] Failed to inject heartbeat script (serve_react): {e}")
resp = send_from_directory(app.static_folder, 'index.html')
resp.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
resp.headers['Pragma'] = 'no-cache'
resp.headers['Expires'] = '0'
return resp
resp = make_response(html)
resp.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
resp.headers['Pragma'] = 'no-cache'
resp.headers['Expires'] = '0'
return resp
@app.route('/similar-images', methods=['GET'])
@fresh_login_required
def similar_images_page():
"""Serve similar images search page"""
resp = send_from_directory(app.static_folder, 'similar-images.html')
resp.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
resp.headers['Pragma'] = 'no-cache'
resp.headers['Expires'] = '0'
return resp
@app.route('/object-detection-search', methods=['GET'])
@fresh_login_required
def object_detection_search_page():
"""Serve object detection search page"""
resp = send_from_directory(app.static_folder, 'object-detection-search.html')
resp.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
resp.headers['Pragma'] = 'no-cache'
resp.headers['Expires'] = '0'
return resp
@app.route('/model-vector-db', methods=['GET'])
@fresh_login_required
def model_vector_db_page():
"""Serve model vector DB UI page"""
resp = send_from_directory(app.static_folder, 'model-vector-db.html')
resp.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
resp.headers['Pragma'] = 'no-cache'
resp.headers['Expires'] = '0'
return resp
@app.route('/openai-chat', methods=['GET'])
@fresh_login_required
def openai_chat_page():
"""Serve OpenAI chat UI page"""
resp = send_from_directory(app.static_folder, 'openai-chat.html')
resp.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
resp.headers['Pragma'] = 'no-cache'
resp.headers['Expires'] = '0'
return resp
@app.route('/api/openai/chat', methods=['POST'])
@fresh_login_required
def openai_chat_api():
"""Forward chat request to OpenAI Chat Completions API.
Expects JSON: { prompt: string, model?: string, api_key?: string, system?: string }
Uses OPENAI_API_KEY from environment if api_key not provided.
"""
try:
data = request.get_json(force=True)
except Exception:
return jsonify({"error": "Invalid JSON body"}), 400
prompt = (data or {}).get('prompt', '').strip()
model = (data or {}).get('model') or os.environ.get('OPENAI_MODEL', 'gpt-4')
system = (data or {}).get('system') or 'You are a helpful assistant.'
api_key = (data or {}).get('api_key') or os.environ.get('OPENAI_API_KEY')
if not prompt:
return jsonify({"error": "Missing 'prompt'"}), 400
if not api_key:
return jsonify({"error": "Missing OpenAI API key. Provide in request or set OPENAI_API_KEY env."}), 400
# Prefer official Python SDK if available
if OpenAI is None:
return jsonify({"error": "OpenAI Python package not installed on server"}), 500
try:
start = time.time()
client = OpenAI(api_key=api_key)
# Perform chat completion
chat = client.chat.completions.create(
model=model,
messages=[
{"role": "system", "content": system},
{"role": "user", "content": prompt},
],
)
latency = round(time.time() - start, 3)
except Exception as e:
# Log detailed error for diagnostics
try:
import traceback
traceback.print_exc()
except Exception:
pass
err_msg = str(e)
print(f"[OpenAI Chat Error] model={model} err={err_msg}")
return jsonify({
"error": "OpenAI SDK call failed",
"detail": err_msg,
"model": model
}), 502
try:
content = chat.choices[0].message.content if chat and chat.choices else ''
usage = getattr(chat, 'usage', None)
usage = usage.model_dump() if hasattr(usage, 'model_dump') else (usage or {})
except Exception as e:
return jsonify({"error": f"Failed to parse SDK response: {str(e)}"}), 500
return jsonify({
'response': content,
'model': model,
'usage': usage,
'latency_sec': latency
})
@app.route('/api/vision-rag/query', methods=['POST'])
@login_required
def vision_rag_query():
"""Vision RAG endpoint.
Expects JSON with one of the following query modes and a user question:
- { userQuery, searchType: 'image', image, n_results? }
- { userQuery, searchType: 'object', objectId, n_results? }
- { userQuery, searchType: 'class', class_name, n_results? }
Returns: { answer, retrieved: [...], model, latency_sec }
"""
if ChatOpenAI is None:
return jsonify({"error": "LangChain not installed on server"}), 500
data = request.get_json(silent=True) or {}
user_query = (data.get('userQuery') or '').strip()
if not user_query:
return jsonify({"error": "Missing 'userQuery'"}), 400
api_key = data.get('api_key') or os.environ.get('OPENAI_API_KEY')
if not api_key:
return jsonify({"error": "Missing OpenAI API key. Provide in request or set OPENAI_API_KEY env."}), 400
search_type = data.get('searchType', 'image')
n_results = int(data.get('n_results', 5))
# Build query embedding or filtered fetch similar to /api/search-similar-objects
results = None
try:
if search_type == 'image' and 'image' in data:
image_data = data['image']
if isinstance(image_data, str) and image_data.startswith('data:image'):
image_data = image_data.split(',')[1]
image = Image.open(BytesIO(base64.b64decode(image_data))).convert('RGB')
query_embedding = generate_image_embedding(image)
if query_embedding is None:
return jsonify({"error": "Failed to generate image embedding"}), 500
results = object_collection.query(
query_embeddings=[query_embedding],
n_results=n_results,
include=["metadatas", "distances"]
) if object_collection is not None else None
elif search_type == 'object' and 'objectId' in data:
obj_id = data['objectId']
base = object_collection.get(ids=[obj_id], include=["embeddings"]) if object_collection is not None else None
emb = base["embeddings"][0] if base and "embeddings" in base and base["embeddings"] else None
if emb is None:
return jsonify({"error": "objectId not found or has no embedding"}), 400
results = object_collection.query(
query_embeddings=[emb],
n_results=n_results,
include=["metadatas", "distances"]
)
elif search_type == 'class' and 'class_name' in data:
filter_query = {"class": {"$eq": data['class_name']}}
results = object_collection.get(
where=filter_query,
limit=n_results,
include=["metadatas", "embeddings", "documents"]
) if object_collection is not None else None
else:
return jsonify({"error": "Invalid search parameters"}), 400
except Exception as e:
return jsonify({"error": f"Retrieval failed: {str(e)}"}), 500
# Format results using existing helper
formatted = format_object_results(results) if results else []
# Build concise context for LLM
def _shorten(md):
try:
bbox = md.get('bbox') if isinstance(md, dict) else None
if isinstance(bbox, dict):
bbox = {k: round(float(v), 3) for k, v in bbox.items() if isinstance(v, (int, float))}
return {
'image_id': md.get('image_id'),
'class': md.get('class'),
'confidence': md.get('confidence'),
'bbox': bbox,
}
except Exception:
return {k: md.get(k) for k in ('image_id', 'class', 'confidence') if k in md}
context_items = []
for r in formatted[:n_results]:
md = r.get('metadata', {})
item = {
'id': r.get('id'),
'distance': r.get('distance'),
'meta': _shorten(md)
}
context_items.append(item)
# Compose prompt
system_text = (
"You are a vision assistant. Use ONLY the provided detected object context to answer. "
"Be concise and state uncertainty if context is insufficient."
)
# Provide the minimal JSON-like context to the model
context_text = json.dumps(context_items, ensure_ascii=False, indent=2)
user_text = f"User question: {user_query}\n\nDetected context (top {len(context_items)}):\n{context_text}"
try:
start = time.time()
llm = ChatOpenAI(api_key=api_key, model=os.environ.get('OPENAI_MODEL', 'gpt-4o'))
# Keep it simple: template -> LLM -> string
prompt = ChatPromptTemplate.from_messages([
("system", system_text),
("human", "{input}")
])
chain = prompt | llm | StrOutputParser()
answer = chain.invoke({"input": user_text})
latency = round(time.time() - start, 3)
except Exception as e:
return jsonify({"error": f"LLM call failed: {str(e)}"}), 502
return jsonify({
"answer": answer,
"retrieved": context_items,
"model": getattr(llm, 'model', None),
"latency_sec": latency
})
@app.route('/api/status', methods=['GET'])
@fresh_login_required
def status():
return jsonify({
"status": "online",
"models": {
"yolo": yolo_model is not None,
"detr": detr_model is not None and detr_processor is not None,
"vit": vit_model is not None and vit_processor is not None
},
"device": "GPU" if torch.cuda.is_available() else "CPU",
"user": current_user.username
})
# Root route is now handled by serve_react function
# This route is removed to prevent conflicts
@app.route('/index')
@login_required
def index_page():
# /index ๊ฒฝ๋กœ๋Š” index.html๋กœ ๋ฆฌ๋””๋ ‰์…˜
print("Index route redirecting to index.html")
return redirect('/index.html')
if __name__ == "__main__":
# ํ—ˆ๊น…ํŽ˜์ด์Šค Space์—์„œ๋Š” PORT ํ™˜๊ฒฝ ๋ณ€์ˆ˜๋ฅผ ์‚ฌ์šฉํ•ฉ๋‹ˆ๋‹ค
port = int(os.environ.get("PORT", 7860))
app.run(debug=False, host='0.0.0.0', port=port)