Spaces:
Running
on
Zero
Running
on
Zero
#!/usr/bin/env python3 | |
import os | |
import re | |
import tempfile | |
import gc | |
from collections.abc import Iterator | |
from threading import Thread | |
import json | |
import requests | |
import cv2 | |
import gradio as gr | |
import spaces | |
import torch | |
import numpy as np | |
from loguru import logger | |
from PIL import Image | |
import time | |
import warnings | |
from typing import Dict, List, Optional, Union | |
import base64 | |
from io import BytesIO | |
# llama-cpp-python for GGUF | |
from llama_cpp import Llama | |
from llama_cpp.llama_chat_format import Llava16ChatHandler | |
# Model download | |
from huggingface_hub import hf_hub_download | |
# CSV/TXT 분석 | |
import pandas as pd | |
# PDF 텍스트 추출 | |
import PyPDF2 | |
warnings.filterwarnings('ignore') | |
print("🎮 로봇 시각 시스템 초기화 (Gemma3-4B GGUF Q4_K_M)...") | |
############################################################################## | |
# 상수 정의 | |
############################################################################## | |
MAX_CONTENT_CHARS = 2000 | |
MAX_INPUT_LENGTH = 2096 | |
MAX_NUM_IMAGES = 5 | |
SERPHOUSE_API_KEY = os.getenv("SERPHOUSE_API_KEY", "") | |
############################################################################## | |
# 전역 변수 | |
############################################################################## | |
llm = None | |
model_loaded = False | |
model_name = "Gemma3-4B-GGUF-Q4_K_M" | |
############################################################################## | |
# 메모리 관리 | |
############################################################################## | |
def clear_cuda_cache(): | |
"""CUDA 캐시를 명시적으로 비웁니다.""" | |
if torch.cuda.is_available(): | |
torch.cuda.empty_cache() | |
gc.collect() | |
############################################################################## | |
# 키워드 추출 함수 | |
############################################################################## | |
def extract_keywords(text: str, top_k: int = 5) -> str: | |
"""키워드 추출""" | |
text = re.sub(r"[^a-zA-Z0-9가-힣\s]", "", text) | |
tokens = text.split() | |
seen = set() | |
unique_tokens = [] | |
for token in tokens: | |
if token not in seen and len(token) > 1: | |
seen.add(token) | |
unique_tokens.append(token) | |
key_tokens = unique_tokens[:top_k] | |
return " ".join(key_tokens) | |
############################################################################## | |
# 웹 검색 함수 | |
############################################################################## | |
def do_web_search(query: str) -> str: | |
"""SerpHouse API를 사용한 웹 검색""" | |
try: | |
url = "https://api.serphouse.com/serp/live" | |
params = { | |
"q": query, | |
"domain": "google.com", | |
"serp_type": "web", | |
"device": "desktop", | |
"lang": "ko", | |
"num": "10" | |
} | |
headers = { | |
"Authorization": f"Bearer {SERPHOUSE_API_KEY}" | |
} | |
logger.info(f"웹 검색 중... 검색어: {query}") | |
response = requests.get(url, headers=headers, params=params, timeout=60) | |
response.raise_for_status() | |
data = response.json() | |
results = data.get("results", {}) | |
organic = results.get("organic", []) if isinstance(results, dict) else [] | |
if not organic: | |
return "검색 결과를 찾을 수 없습니다." | |
max_results = min(10, len(organic)) | |
limited_organic = organic[:max_results] | |
summary_lines = [] | |
for idx, item in enumerate(limited_organic, start=1): | |
title = item.get("title", "제목 없음") | |
link = item.get("link", "#") | |
snippet = item.get("snippet", "설명 없음") | |
displayed_link = item.get("displayed_link", link) | |
summary_lines.append( | |
f"### 결과 {idx}: {title}\n\n" | |
f"{snippet}\n\n" | |
f"**출처**: [{displayed_link}]({link})\n\n" | |
f"---\n" | |
) | |
instructions = """# 웹 검색 결과 | |
아래는 검색 결과입니다. 답변 시 이 정보를 활용하세요: | |
1. 각 결과의 제목, 내용, 출처 링크를 참조하세요 | |
2. 관련 출처를 명시적으로 인용하세요 | |
3. 여러 출처의 정보를 종합하여 답변하세요 | |
""" | |
search_results = instructions + "\n".join(summary_lines) | |
return search_results | |
except Exception as e: | |
logger.error(f"웹 검색 실패: {e}") | |
return f"웹 검색 실패: {str(e)}" | |
############################################################################## | |
# 문서 처리 함수 | |
############################################################################## | |
def analyze_csv_file(path: str) -> str: | |
"""CSV 파일 분석""" | |
try: | |
df = pd.read_csv(path) | |
if df.shape[0] > 50 or df.shape[1] > 10: | |
df = df.iloc[:50, :10] | |
df_str = df.to_string() | |
if len(df_str) > MAX_CONTENT_CHARS: | |
df_str = df_str[:MAX_CONTENT_CHARS] + "\n...(중략)..." | |
return f"**[CSV 파일: {os.path.basename(path)}]**\n\n{df_str}" | |
except Exception as e: | |
return f"CSV 읽기 실패 ({os.path.basename(path)}): {str(e)}" | |
def analyze_txt_file(path: str) -> str: | |
"""TXT 파일 분석""" | |
try: | |
with open(path, "r", encoding="utf-8") as f: | |
text = f.read() | |
if len(text) > MAX_CONTENT_CHARS: | |
text = text[:MAX_CONTENT_CHARS] + "\n...(중략)..." | |
return f"**[TXT 파일: {os.path.basename(path)}]**\n\n{text}" | |
except Exception as e: | |
return f"TXT 읽기 실패 ({os.path.basename(path)}): {str(e)}" | |
def pdf_to_markdown(pdf_path: str) -> str: | |
"""PDF를 마크다운으로 변환""" | |
text_chunks = [] | |
try: | |
with open(pdf_path, "rb") as f: | |
reader = PyPDF2.PdfReader(f) | |
max_pages = min(5, len(reader.pages)) | |
for page_num in range(max_pages): | |
page = reader.pages[page_num] | |
page_text = page.extract_text() or "" | |
page_text = page_text.strip() | |
if page_text: | |
if len(page_text) > MAX_CONTENT_CHARS // max_pages: | |
page_text = page_text[:MAX_CONTENT_CHARS // max_pages] + "...(중략)" | |
text_chunks.append(f"## 페이지 {page_num+1}\n\n{page_text}\n") | |
if len(reader.pages) > max_pages: | |
text_chunks.append(f"\n...({max_pages}/{len(reader.pages)} 페이지 표시)...") | |
except Exception as e: | |
return f"PDF 읽기 실패 ({os.path.basename(pdf_path)}): {str(e)}" | |
full_text = "\n".join(text_chunks) | |
if len(full_text) > MAX_CONTENT_CHARS: | |
full_text = full_text[:MAX_CONTENT_CHARS] + "\n...(중략)..." | |
return f"**[PDF 파일: {os.path.basename(pdf_path)}]**\n\n{full_text}" | |
############################################################################## | |
# 이미지를 base64로 변환 | |
############################################################################## | |
def image_to_base64_data_uri(image: Union[np.ndarray, Image.Image]) -> str: | |
"""이미지를 base64 data URI로 변환""" | |
if isinstance(image, np.ndarray): | |
image = Image.fromarray(image).convert('RGB') | |
buffered = BytesIO() | |
image.save(buffered, format="JPEG", quality=85) | |
img_str = base64.b64encode(buffered.getvalue()).decode() | |
return f"data:image/jpeg;base64,{img_str}" | |
############################################################################## | |
# 모델 로드 | |
############################################################################## | |
def download_model_files(): | |
"""Hugging Face Hub에서 모델 파일 다운로드""" | |
# 여러 가능한 저장소 시도 | |
model_repos = [ | |
# 첫 번째 시도: 일반적인 Gemma 3 4B GGUF | |
{ | |
"repo": "Mungert/gemma-3-4b-it-gguf", | |
"model": "google_gemma-3-4b-it-q4_k_m.gguf", | |
"mmproj": "google_gemma-3-4b-it-mmproj-bf16.gguf" | |
}, | |
# 두 번째 시도: LM Studio 버전 | |
{ | |
"repo": "lmstudio-community/gemma-3-4b-it-GGUF", | |
"model": "gemma-3-4b-it-Q4_K_M.gguf", | |
"mmproj": "gemma-3-4b-it-mmproj-f16.gguf" | |
}, | |
# 세 번째 시도: unsloth 버전 | |
{ | |
"repo": "unsloth/gemma-3-4b-it-GGUF", | |
"model": "gemma-3-4b-it.Q4_K_M.gguf", | |
"mmproj": "gemma-3-4b-it.mmproj.gguf" | |
} | |
] | |
for repo_info in model_repos: | |
try: | |
logger.info(f"저장소 시도: {repo_info['repo']}") | |
# 메인 모델 다운로드 | |
model_filename = repo_info["model"] | |
logger.info(f"모델 다운로드 중: {model_filename}") | |
model_path = hf_hub_download( | |
repo_id=repo_info["repo"], | |
filename=model_filename, | |
resume_download=True, | |
local_files_only=False | |
) | |
# Vision projection 파일 다운로드 | |
mmproj_filename = repo_info["mmproj"] | |
logger.info(f"Vision 모델 다운로드 중: {mmproj_filename}") | |
try: | |
mmproj_path = hf_hub_download( | |
repo_id=repo_info["repo"], | |
filename=mmproj_filename, | |
resume_download=True, | |
local_files_only=False | |
) | |
except: | |
# mmproj 파일이 없을 수도 있음 | |
logger.warning(f"Vision 모델을 찾을 수 없습니다: {mmproj_filename}") | |
logger.warning("텍스트 전용 모드로 진행합니다.") | |
mmproj_path = None | |
logger.info(f"✅ 모델 다운로드 성공!") | |
logger.info(f"모델 경로: {model_path}") | |
if mmproj_path: | |
logger.info(f"Vision 경로: {mmproj_path}") | |
return model_path, mmproj_path | |
except Exception as e: | |
logger.error(f"저장소 {repo_info['repo']} 시도 실패: {e}") | |
continue | |
# 모든 시도가 실패한 경우 | |
raise Exception("사용 가능한 GGUF 모델을 찾을 수 없습니다. 인터넷 연결을 확인하세요.") | |
def load_model(): | |
global llm, model_loaded | |
if model_loaded: | |
logger.info("모델이 이미 로드되어 있습니다.") | |
return True | |
try: | |
logger.info("Gemma3-4B GGUF Q4_K_M 모델 로딩 시작...") | |
clear_cuda_cache() | |
# 모델 파일 다운로드 | |
model_path, mmproj_path = download_model_files() | |
# GPU 사용 가능 여부 확인 | |
n_gpu_layers = -1 if torch.cuda.is_available() else 0 | |
# 채팅 핸들러 생성 (비전 지원 - mmproj가 있는 경우만) | |
chat_handler = None | |
if mmproj_path: | |
try: | |
chat_handler = Llava16ChatHandler( | |
clip_model_path=mmproj_path, | |
verbose=False | |
) | |
logger.info("✅ Vision 모델 로드 성공") | |
except Exception as e: | |
logger.warning(f"Vision 모델 로드 실패, 텍스트 전용 모드로 전환: {e}") | |
chat_handler = None | |
# 모델 로드 | |
llm_params = { | |
"model_path": model_path, | |
"n_ctx": 4096, # 컨텍스트 크기 | |
"n_gpu_layers": n_gpu_layers, # GPU 레이어 | |
"n_threads": 8, # CPU 스레드 | |
"verbose": False, | |
"seed": 42, | |
} | |
# chat_handler가 있으면 추가 | |
if chat_handler: | |
llm_params["chat_handler"] = chat_handler | |
llm_params["logits_all"] = True # 비전 모델에 필요 | |
llm = Llama(**llm_params) | |
model_loaded = True | |
logger.info(f"✅ Gemma3-4B 모델 로딩 완료!") | |
if not chat_handler: | |
logger.warning("⚠️ 텍스트 전용 모드로 실행 중입니다. 이미지 분석이 제한될 수 있습니다.") | |
return True | |
except Exception as e: | |
logger.error(f"모델 로딩 실패: {e}") | |
import traceback | |
logger.error(traceback.format_exc()) | |
return False | |
############################################################################## | |
# 채팅 템플릿 포맷팅 | |
############################################################################## | |
def format_chat_prompt(system_prompt: str, user_prompt: str, image_uri: Optional[str] = None) -> List[Dict]: | |
"""Gemma 스타일 채팅 프롬프트 생성""" | |
messages = [] | |
# 시스템 메시지 | |
messages.append({ | |
"role": "system", | |
"content": system_prompt | |
}) | |
# 사용자 메시지 | |
user_content = [] | |
if image_uri: | |
user_content.append({ | |
"type": "image_url", | |
"image_url": {"url": image_uri} | |
}) | |
user_content.append({ | |
"type": "text", | |
"text": user_prompt | |
}) | |
messages.append({ | |
"role": "user", | |
"content": user_content | |
}) | |
return messages | |
############################################################################## | |
# 이미지 분석 (로봇 태스크 중심) | |
############################################################################## | |
def analyze_image_for_robot( | |
image: Union[np.ndarray, Image.Image], | |
prompt: str, | |
task_type: str = "general", | |
use_web_search: bool = False, | |
enable_thinking: bool = False, | |
max_new_tokens: int = 300 | |
) -> str: | |
"""로봇 작업을 위한 이미지 분석""" | |
global llm | |
if not model_loaded: | |
if not load_model(): | |
return "❌ 모델 로딩 실패" | |
try: | |
# Vision 모델이 없는 경우 경고 | |
if not hasattr(llm, 'chat_handler') or llm.chat_handler is None: | |
logger.warning("Vision 모델이 로드되지 않았습니다. 텍스트 기반 분석만 가능합니다.") | |
# 텍스트 전용 분석 | |
system_prompt = f"""당신은 로봇 시각 시스템 시뮬레이터입니다. | |
실제 이미지를 볼 수는 없지만, 사용자의 설명을 바탕으로 로봇 작업을 계획하고 분석합니다. | |
태스크 유형: {task_type}""" | |
messages = [ | |
{"role": "system", "content": system_prompt}, | |
{"role": "user", "content": f"[이미지 분석 요청] {prompt}"} | |
] | |
response = llm.create_chat_completion( | |
messages=messages, | |
max_tokens=max_new_tokens, | |
temperature=0.7, | |
top_p=0.9, | |
stream=False | |
) | |
result = response['choices'][0]['message']['content'].strip() | |
return f"⚠️ 텍스트 전용 모드\n\n{result}" | |
# 이미지를 base64로 변환 | |
image_uri = image_to_base64_data_uri(image) | |
# 태스크별 시스템 프롬프트 구성 | |
system_prompts = { | |
"general": "당신은 로봇 시각 시스템입니다. 먼저 장면을 1-2줄로 설명하고, 핵심 내용을 간결하게 분석하세요.", | |
"planning": """당신은 로봇 작업 계획 AI입니다. | |
먼저 장면 이해를 1-2줄로 설명하고, 그 다음 작업 계획을 작성하세요. | |
형식: | |
[장면 이해] 현재 보이는 장면을 1-2줄로 설명 | |
[작업 계획] | |
Step_1: xxx | |
Step_2: xxx | |
Step_n: xxx""", | |
"grounding": "당신은 객체 위치 시스템입니다. 먼저 보이는 객체들을 한 줄로 설명하고, 요청된 객체 위치를 [x1, y1, x2, y2]로 반환하세요.", | |
"affordance": "당신은 파지점 분석 AI입니다. 먼저 대상 객체를 한 줄로 설명하고, 파지 영역을 [x1, y1, x2, y2]로 반환하세요.", | |
"trajectory": "당신은 경로 계획 AI입니다. 먼저 환경을 한 줄로 설명하고, 경로를 [(x1,y1), (x2,y2), ...]로 제시하세요.", | |
"pointing": "당신은 지점 지정 시스템입니다. 먼저 참조점들을 한 줄로 설명하고, 위치를 [(x1,y1), (x2,y2), ...]로 반환하세요." | |
} | |
system_prompt = system_prompts.get(task_type, system_prompts["general"]) | |
# Chain-of-Thought 추가 (선택적) | |
if enable_thinking: | |
system_prompt += "\n\n추론 과정을 <thinking></thinking> 태그 안에 작성 후 최종 답변을 제시하세요. 장면 이해는 추론 과정과 별도로 반드시 포함하세요." | |
# 웹 검색 수행 | |
combined_system = system_prompt | |
if use_web_search: | |
keywords = extract_keywords(prompt, top_k=5) | |
if keywords: | |
logger.info(f"웹 검색 키워드: {keywords}") | |
search_results = do_web_search(keywords) | |
combined_system = f"{search_results}\n\n{system_prompt}" | |
# 메시지 구성 | |
messages = format_chat_prompt(combined_system, prompt, image_uri) | |
# 생성 | |
response = llm.create_chat_completion( | |
messages=messages, | |
max_tokens=max_new_tokens, | |
temperature=0.7, | |
top_p=0.9, | |
stream=False | |
) | |
# 응답 추출 | |
result = response['choices'][0]['message']['content'].strip() | |
return result | |
except Exception as e: | |
logger.error(f"이미지 분석 오류: {e}") | |
import traceback | |
return f"❌ 분석 오류: {str(e)}\n{traceback.format_exc()}" | |
finally: | |
clear_cuda_cache() | |
############################################################################## | |
# 문서 분석 (스트리밍) | |
############################################################################## | |
def analyze_documents_streaming( | |
files: List[str], | |
prompt: str, | |
use_web_search: bool = False, | |
max_new_tokens: int = 2048 | |
) -> Iterator[str]: | |
"""문서 분석 (스트리밍)""" | |
global llm | |
if not model_loaded: | |
if not load_model(): | |
yield "❌ 모델 로딩 실패" | |
return | |
try: | |
# 시스템 프롬프트 | |
system_content = "당신은 문서를 분석하고 요약하는 전문 AI입니다." | |
# 웹 검색 | |
if use_web_search: | |
keywords = extract_keywords(prompt, top_k=5) | |
if keywords: | |
search_results = do_web_search(keywords) | |
system_content = f"{search_results}\n\n{system_content}" | |
# 문서 내용 처리 | |
doc_contents = [] | |
for file_path in files: | |
if file_path.lower().endswith('.csv'): | |
content = analyze_csv_file(file_path) | |
elif file_path.lower().endswith('.txt'): | |
content = analyze_txt_file(file_path) | |
elif file_path.lower().endswith('.pdf'): | |
content = pdf_to_markdown(file_path) | |
else: | |
continue | |
doc_contents.append(content) | |
# 전체 프롬프트 구성 | |
full_prompt = "\n\n".join(doc_contents) + f"\n\n{prompt}" | |
# 메시지 구성 | |
messages = [ | |
{"role": "system", "content": system_content}, | |
{"role": "user", "content": full_prompt} | |
] | |
# 스트리밍 생성 | |
stream = llm.create_chat_completion( | |
messages=messages, | |
max_tokens=max_new_tokens, | |
temperature=0.8, | |
top_p=0.9, | |
stream=True | |
) | |
# 스트리밍 출력 | |
output = "" | |
for chunk in stream: | |
if 'choices' in chunk and len(chunk['choices']) > 0: | |
delta = chunk['choices'][0].get('delta', {}) | |
if 'content' in delta: | |
output += delta['content'] | |
yield output | |
except Exception as e: | |
logger.error(f"문서 분석 오류: {e}") | |
yield f"❌ 오류 발생: {str(e)}" | |
finally: | |
clear_cuda_cache() | |
############################################################################## | |
# Gradio UI (로봇 시각화 중심) | |
############################################################################## | |
css = """ | |
.robot-header { | |
text-align: center; | |
background: linear-gradient(135deg, #1e3c72 0%, #2a5298 50%, #667eea 100%); | |
color: white; | |
padding: 20px; | |
border-radius: 10px; | |
margin-bottom: 20px; | |
box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); | |
} | |
.status-box { | |
text-align: center; | |
padding: 10px; | |
border-radius: 5px; | |
margin: 10px 0; | |
font-weight: bold; | |
} | |
.info-box { | |
background: #f0f0f0; | |
padding: 15px; | |
border-radius: 8px; | |
margin: 10px 0; | |
border-left: 4px solid #2a5298; | |
} | |
.task-button { | |
min-height: 60px; | |
font-size: 1.1em; | |
} | |
.webcam-container { | |
border: 3px solid #2a5298; | |
border-radius: 10px; | |
padding: 10px; | |
background: #f8f9fa; | |
} | |
.auto-capture-status { | |
text-align: center; | |
padding: 5px; | |
border-radius: 5px; | |
margin: 5px 0; | |
font-weight: bold; | |
background: #e8f5e9; | |
color: #2e7d32; | |
} | |
.model-info { | |
background: #fff3cd; | |
color: #856404; | |
padding: 10px; | |
border-radius: 5px; | |
margin: 10px 0; | |
text-align: center; | |
} | |
""" | |
with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B GGUF)", css=css) as demo: | |
gr.HTML(""" | |
<div class="robot-header"> | |
<h1>🤖 로봇 시각 시스템</h1> | |
<h3>🎮 Gemma3-4B GGUF Q4_K_M + 📷 실시간 웹캠 + 🔍 웹 검색</h3> | |
<p>⚡ 양자화 모델로 더 빠르고 효율적인 로봇 작업 분석!</p> | |
</div> | |
""") | |
gr.HTML(""" | |
<div class="model-info"> | |
<strong>모델:</strong> Gemma3-4B Q4_K_M (2.5GB) | <strong>메모리 사용:</strong> ~3-4GB VRAM | |
</div> | |
""") | |
with gr.Row(): | |
# 왼쪽: 웹캠 및 입력 | |
with gr.Column(scale=1): | |
gr.Markdown("### 📷 실시간 웹캠") | |
with gr.Group(elem_classes="webcam-container"): | |
webcam = gr.Image( | |
sources=["webcam"], | |
streaming=True, | |
type="numpy", | |
label="실시간 스트리밍", | |
height=350 | |
) | |
# 자동 캡처 상태 표시 | |
auto_capture_status = gr.HTML( | |
'<div class="auto-capture-status">🔄 자동 캡처: 대기 중</div>' | |
) | |
# 캡처된 이미지 표시 | |
captured_image = gr.Image( | |
label="캡처된 이미지", | |
height=200, | |
visible=False | |
) | |
# 로봇 작업 버튼들 | |
gr.Markdown("### 🎯 로봇 작업 선택") | |
with gr.Row(): | |
capture_btn = gr.Button("📸 수동 캡처", variant="primary", elem_classes="task-button") | |
clear_capture_btn = gr.Button("🗑️ 초기화", elem_classes="task-button") | |
with gr.Row(): | |
auto_capture_toggle = gr.Checkbox( | |
label="🔄 자동 캡처 활성화 (10초마다)", | |
value=False, | |
info="활성화 시 10초마다 자동으로 캡처 및 분석" | |
) | |
with gr.Row(): | |
planning_btn = gr.Button("📋 작업 계획", elem_classes="task-button") | |
grounding_btn = gr.Button("📍 객체 위치", elem_classes="task-button") | |
with gr.Row(): | |
affordance_btn = gr.Button("🤏 파지점 분석", elem_classes="task-button") | |
trajectory_btn = gr.Button("🛤️ 경로 계획", elem_classes="task-button") | |
# 오른쪽: 분석 설정 및 결과 | |
with gr.Column(scale=2): | |
gr.Markdown("### ⚙️ 분석 설정") | |
with gr.Row(): | |
with gr.Column(): | |
task_prompt = gr.Textbox( | |
label="작업 설명 / 질문", | |
placeholder="예: 테이블 위의 컵을 잡아서 싱크대에 놓기", | |
value="현재 장면을 분석하고 로봇이 수행할 수 있는 작업을 제안하세요.", | |
lines=2 | |
) | |
with gr.Row(): | |
use_web_search = gr.Checkbox( | |
label="🔍 웹 검색 사용", | |
value=False, | |
info="관련 정보를 웹에서 검색합니다" | |
) | |
enable_thinking = gr.Checkbox( | |
label="🤔 추론 과정 표시", | |
value=False, | |
info="Chain-of-Thought 추론 과정을 보여줍니다" | |
) | |
max_tokens = gr.Slider( | |
label="최대 토큰 수", | |
minimum=100, | |
maximum=2048, | |
value=300, | |
step=50 | |
) | |
gr.Markdown("### 📊 분석 결과") | |
result_output = gr.Textbox( | |
label="AI 분석 결과", | |
lines=20, | |
max_lines=40, | |
show_copy_button=True, | |
elem_id="result" | |
) | |
status_display = gr.HTML( | |
'<div class="status-box" style="background:#d4edda; color:#155724;">🎮 시스템 준비 완료</div>' | |
) | |
# 문서 분석 탭 | |
with gr.Tab("📄 문서 분석", visible=False): | |
with gr.Row(): | |
with gr.Column(): | |
doc_files = gr.File( | |
label="문서 업로드", | |
file_count="multiple", | |
file_types=[".pdf", ".csv", ".txt"], | |
type="filepath" | |
) | |
doc_prompt = gr.Textbox( | |
label="분석 요청", | |
placeholder="예: 이 문서들의 핵심 내용을 요약하고 비교 분석하세요.", | |
lines=3 | |
) | |
doc_web_search = gr.Checkbox( | |
label="🔍 웹 검색 사용", | |
value=False | |
) | |
analyze_docs_btn = gr.Button("📊 문서 분석", variant="primary") | |
with gr.Column(): | |
doc_result = gr.Textbox( | |
label="분석 결과", | |
lines=25, | |
max_lines=50 | |
) | |
# 이벤트 핸들러 | |
webcam_state = gr.State(None) | |
auto_capture_state = gr.State({"enabled": False, "timer": None}) | |
def capture_webcam(frame): | |
"""웹캠 프레임 캡처""" | |
if frame is None: | |
return None, None, '<div class="status-box" style="background:#f8d7da; color:#721c24;">❌ 웹캠 프레임 없음</div>' | |
return frame, gr.update(value=frame, visible=True), '<div class="status-box" style="background:#d4edda; color:#155724;">✅ 이미지 캡처 완료</div>' | |
def clear_capture(): | |
"""캡처 초기화""" | |
return None, gr.update(visible=False), '<div class="status-box" style="background:#d4edda; color:#155724;">🎮 시스템 준비 완료</div>' | |
def analyze_with_task(image, prompt, task_type, use_search, thinking, tokens): | |
"""특정 태스크로 이미지 분석""" | |
if image is None: | |
return "❌ 먼저 이미지를 캡처하세요.", '<div class="status-box" style="background:#f8d7da; color:#721c24;">❌ 이미지 없음</div>' | |
status = f'<div class="status-box" style="background:#cce5ff; color:#004085;">🚀 {task_type} 분석 중...</div>' | |
result = analyze_image_for_robot( | |
image=image, | |
prompt=prompt, | |
task_type=task_type, | |
use_web_search=use_search, | |
enable_thinking=thinking, | |
max_new_tokens=tokens | |
) | |
# 결과 포맷팅 | |
timestamp = time.strftime("%H:%M:%S") | |
task_names = { | |
"planning": "작업 계획", | |
"grounding": "객체 위치", | |
"affordance": "파지점", | |
"trajectory": "경로 계획" | |
} | |
formatted_result = f"""🤖 {task_names.get(task_type, '분석')} 결과 ({timestamp}) | |
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ | |
{result} | |
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━""" | |
complete_status = '<div class="status-box" style="background:#d4edda; color:#155724;">✅ 분석 완료!</div>' | |
return formatted_result, complete_status | |
# 자동 캡처 및 분석 함수 | |
def auto_capture_and_analyze(webcam_frame, task_prompt, use_search, thinking, tokens, auto_state): | |
"""자동 캡처 및 분석""" | |
if webcam_frame is None: | |
return ( | |
None, | |
"자동 캡처 대기 중...", | |
'<div class="status-box" style="background:#fff3cd; color:#856404;">⏳ 웹캠 대기 중</div>', | |
'<div class="auto-capture-status">🔄 자동 캡처: 웹캠 대기 중</div>' | |
) | |
# 캡처 수행 | |
timestamp = time.strftime("%H:%M:%S") | |
# 이미지 분석 (작업 계획 모드로) | |
result = analyze_image_for_robot( | |
image=webcam_frame, | |
prompt=task_prompt, | |
task_type="planning", | |
use_web_search=use_search, | |
enable_thinking=thinking, | |
max_new_tokens=tokens | |
) | |
formatted_result = f"""🔄 자동 분석 완료 ({timestamp}) | |
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ | |
{result} | |
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━""" | |
return ( | |
webcam_frame, | |
formatted_result, | |
'<div class="status-box" style="background:#d4edda; color:#155724;">✅ 자동 분석 완료</div>', | |
f'<div class="auto-capture-status">🔄 자동 캡처: 마지막 분석 {timestamp}</div>' | |
) | |
# 웹캠 스트리밍 | |
webcam.stream( | |
fn=lambda x: x, | |
inputs=[webcam], | |
outputs=[webcam_state] | |
) | |
# 수동 캡처 버튼 | |
capture_btn.click( | |
fn=capture_webcam, | |
inputs=[webcam_state], | |
outputs=[webcam_state, captured_image, status_display] | |
) | |
# 초기화 버튼 | |
clear_capture_btn.click( | |
fn=clear_capture, | |
outputs=[webcam_state, captured_image, status_display] | |
) | |
# 작업 버튼들 | |
planning_btn.click( | |
fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "planning", s, t, tk), | |
inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens], | |
outputs=[result_output, status_display] | |
) | |
grounding_btn.click( | |
fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "grounding", s, t, tk), | |
inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens], | |
outputs=[result_output, status_display] | |
) | |
affordance_btn.click( | |
fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "affordance", s, t, tk), | |
inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens], | |
outputs=[result_output, status_display] | |
) | |
trajectory_btn.click( | |
fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "trajectory", s, t, tk), | |
inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens], | |
outputs=[result_output, status_display] | |
) | |
# 문서 분석 | |
def analyze_docs(files, prompt, use_search): | |
if not files: | |
return "❌ 문서를 업로드하세요." | |
output = "" | |
for chunk in analyze_documents_streaming(files, prompt, use_search): | |
output = chunk | |
return output | |
analyze_docs_btn.click( | |
fn=analyze_docs, | |
inputs=[doc_files, doc_prompt, doc_web_search], | |
outputs=[doc_result] | |
) | |
# 자동 캡처 타이머 (10초마다) | |
timer = gr.Timer(10.0, active=False) | |
# 자동 캡처 토글 이벤트 | |
def toggle_auto_capture(enabled): | |
if enabled: | |
return gr.Timer(10.0, active=True), '<div class="auto-capture-status">🔄 자동 캡처: 활성화됨 (10초마다)</div>' | |
else: | |
return gr.Timer(active=False), '<div class="auto-capture-status">🔄 자동 캡처: 비활성화됨</div>' | |
auto_capture_toggle.change( | |
fn=toggle_auto_capture, | |
inputs=[auto_capture_toggle], | |
outputs=[timer, auto_capture_status] | |
) | |
# 타이머 틱 이벤트 | |
timer.tick( | |
fn=auto_capture_and_analyze, | |
inputs=[webcam_state, task_prompt, use_web_search, enable_thinking, max_tokens, auto_capture_state], | |
outputs=[captured_image, result_output, status_display, auto_capture_status] | |
) | |
# 초기 모델 로드 | |
def initial_load(): | |
# 첫 실행 시 GPU에서 모델 로드 | |
return "시스템 준비 완료! 첫 분석 시 모델이 자동으로 로드됩니다. 🚀" | |
demo.load( | |
fn=initial_load, | |
outputs=None | |
) | |
if __name__ == "__main__": | |
print("🚀 로봇 시각 시스템 시작 (Gemma3-4B GGUF Q4_K_M)...") | |
demo.launch( | |
server_name="0.0.0.0", | |
server_port=7860, | |
share=False, | |
show_error=True, | |
debug=False | |
) |