Spaces:
Runtime error
Runtime error
| import os | |
| import io | |
| import json | |
| import math | |
| import time | |
| import traceback | |
| import subprocess | |
| import numpy as np | |
| import cv2 | |
| import gradio as gr | |
| import torch | |
| import chromadb | |
| from chromadb.config import Settings | |
| import ollama | |
| from PIL import Image as PILImage | |
| from sentence_transformers import SentenceTransformer | |
| from ultralytics import YOLO | |
| # --------------------------------------------------------- | |
| # Configuration & Paths (Local Structure) | |
| # --------------------------------------------------------- | |
| ARTIFACTS_DIR = os.path.join(os.getcwd(), 'artifacts') | |
| WEIGHTS_PATH = os.path.join(ARTIFACTS_DIR, 'best.pt') | |
| RAG_DB_PATH = os.path.join(os.getcwd(), 'rag_knowledge', 'chroma_db') | |
| CONF_THRESH = 0.35 | |
| IOU_THRESH = 0.35 | |
| IMG_SIZE = 640 | |
| OLLAMA_MODEL = "gemma4:e4b" | |
| OLLAMA_HOST = "http://127.0.0.1:11434" | |
| COCO_NAME_TO_UNIFIED = { | |
| 'person':'person','chair':'chair','couch':'couch','bed':'bed', | |
| 'dining table':'dining table','toilet':'toilet','tv':'tv', | |
| 'laptop':'laptop','microwave':'microwave','oven':'oven', | |
| 'refrigerator':'refrigerator','sink':'sink','clock':'clock', | |
| 'vase':'vase','bottle':'bottle','cup':'cup','book':'book', | |
| 'potted plant':'potted plant', | |
| 'lamp':'lamp', 'door':'door', 'window':'window' | |
| } | |
| COLOR_NAMES = [ | |
| ('red', ( 0, 80, 80), ( 10,255,255)), ('red', (170, 80, 80), (180,255,255)), | |
| ('orange', ( 11, 80, 80), ( 25,255,255)), ('yellow', ( 26, 80, 80), ( 34,255,255)), | |
| ('green', ( 35, 60, 60), ( 85,255,255)), ('blue', ( 95, 60, 60), (130,255,255)), | |
| ('purple', (131, 60, 60), (160,255,255)), ('brown', ( 5, 60, 40), ( 20,180,180)), | |
| ('white', ( 0, 0,200), (180, 30,255)), ('gray', ( 0, 0, 80), (180, 30,200)), | |
| ('black', ( 0, 0, 0), (180,255, 70)), | |
| ] | |
| _state = {'scene_json': '[]', 'rag_ctx': '', 'last_answer': ''} | |
| print('Loading models into memory (this will take a minute)...') | |
| _errors = [] | |
| try: | |
| _yolo = YOLO(WEIGHTS_PATH) | |
| _GPU = '0' if torch.cuda.is_available() else 'cpu' | |
| _yolo.predict(source=np.zeros((640,640,3),dtype=np.uint8), imgsz=IMG_SIZE, conf=CONF_THRESH, device=_GPU, verbose=False) | |
| print(f' ✓ YOLO (device={_GPU})') | |
| except Exception as e: | |
| print(f' ✗ YOLO: {e}\n (Did you place best.pt in {WEIGHTS_PATH}?)') | |
| _errors.append('YOLO'); _yolo = None | |
| try: | |
| _emb = SentenceTransformer('all-MiniLM-L6-v2') | |
| print(' ✓ Embedder') | |
| except Exception as e: | |
| print(f' ✗ Embedder: {e}'); _errors.append('Embedder'); _emb = None | |
| try: | |
| _rag_client = chromadb.PersistentClient(path=RAG_DB_PATH, settings=Settings(anonymized_telemetry=False)) | |
| _rag = _rag_client.get_collection('vqa_knowledge') | |
| print(f' ✓ ChromaDB ({_rag.count()} chunks)') | |
| except Exception as e: | |
| print(f' ✗ ChromaDB: {e}\n (Did you place chroma_db in {RAG_DB_PATH}?)') | |
| _errors.append('ChromaDB'); _rag = None | |
| try: | |
| import requests | |
| try: | |
| requests.get(f"{OLLAMA_HOST}/", timeout=3) | |
| print(' ✓ Ollama server running') | |
| except Exception: | |
| print(' ⚠ Ollama server not responding. Please run `ollama serve` manually.') | |
| _ollama_client = ollama.Client(host=OLLAMA_HOST) | |
| _ollama_client.show(OLLAMA_MODEL) | |
| print(f' ✓ Ollama model ({OLLAMA_MODEL}) ready') | |
| except Exception as e: | |
| print(f' ✗ Ollama: {e}\n (Did you run `ollama pull {OLLAMA_MODEL}`?)') | |
| _errors.append('Ollama'); _ollama_client = None | |
| def get_dominant_color(crop_bgr): | |
| if crop_bgr is None or crop_bgr.size == 0: return 'unknown' | |
| try: | |
| small = cv2.resize(crop_bgr, (32,32), interpolation=cv2.INTER_AREA) | |
| hsv = cv2.cvtColor(small, cv2.COLOR_BGR2HSV) | |
| px = hsv.reshape(-1,3).astype(np.float32) | |
| crit = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0) | |
| _, _, centers = cv2.kmeans(px,1,None,crit,3,cv2.KMEANS_RANDOM_CENTERS) | |
| h,s,v = int(centers[0][0]), int(centers[0][1]), int(centers[0][2]) | |
| for name,(hlo,slo,vlo),(hhi,shi,vhi) in COLOR_NAMES: | |
| if hlo<=h<=hhi and slo<=s<=shi and vlo<=v<=vhi: return name | |
| return 'unknown' | |
| except: return 'unknown' | |
| def _detect(img_np): | |
| if _yolo is None: return [], 'YOLO not loaded.' | |
| gpu = '0' if torch.cuda.is_available() else 'cpu' | |
| img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) | |
| H, W = img_bgr.shape[:2] | |
| r = _yolo.predict(source=img_bgr, imgsz=IMG_SIZE, conf=CONF_THRESH, iou=IOU_THRESH, agnostic_nms=True, device=gpu, verbose=False)[0] | |
| if r.boxes is None or len(r.boxes) == 0: | |
| return [], 'I could not see any objects clearly. Please try a clearer photo.' | |
| dets = [] | |
| for box in r.boxes: | |
| raw_name = _yolo.names[int(box.cls[0])] | |
| if raw_name not in COCO_NAME_TO_UNIFIED: continue | |
| label = COCO_NAME_TO_UNIFIED[raw_name] | |
| conf = float(box.conf[0]) | |
| x1,y1,x2,y2 = map(int, box.xyxy[0].tolist()) | |
| x1,y1 = max(0,x1), max(0,y1) | |
| x2,y2 = min(W,x2), min(H,y2) | |
| crop = img_bgr[y1:y2, x1:x2] | |
| color = get_dominant_color(crop) | |
| xc_n = ((x1+x2)/2)/W; yc_n = ((y1+y2)/2)/H | |
| w_n = (x2-x1)/W; h_n = (y2-y1)/H | |
| pos = 'left' if xc_n<0.33 else ('right' if xc_n>0.67 else 'center') | |
| sz = 'large' if w_n*h_n>0.18 else ('small' if w_n*h_n<0.03 else 'medium') | |
| area = w_n * h_n | |
| dep = 'close' if area > 0.30 else ('far' if area < 0.05 else 'middle') | |
| dets.append({'label':label,'color':color,'size':sz,'position':pos,'depth':dep,'confidence':round(conf,2),'relations':[]}) | |
| seen = {} | |
| for d in sorted(dets, key=lambda x: -x['confidence']): | |
| key = f"{d['label']}_{d['position']}" | |
| if key not in seen: seen[key] = d | |
| return list(seen.values()), None | |
| def _retrieve(query): | |
| if _rag is None or _emb is None: return '' | |
| try: | |
| qv = _emb.encode([query], normalize_embeddings=True)[0].tolist() | |
| res = _rag.query(query_embeddings=[qv], n_results=2, include=['documents']) | |
| return ' '.join(res['documents'][0]) | |
| except: return '' | |
| def _call_gemma_vision(prompt_text, pil_img): | |
| if _ollama_client is None: return 'Ollama not loaded.' | |
| try: | |
| img_small = pil_img.copy() | |
| img_small.thumbnail((768, 768)) | |
| buf = io.BytesIO() | |
| img_small.save(buf, format="JPEG", quality=85) | |
| img_bytes = buf.getvalue() | |
| resp = _ollama_client.chat(model=OLLAMA_MODEL, messages=[{"role": "user", "content": prompt_text, "images": [img_bytes]}], think=False, options={"temperature": 0.3, "num_predict": 120}) | |
| return resp["message"]["content"].strip() | |
| except Exception as e: return f'Ollama error: {e}' | |
| def _call_gemma_text(prompt_text): | |
| if _ollama_client is None: return 'Ollama not loaded.' | |
| try: | |
| resp = _ollama_client.chat(model=OLLAMA_MODEL, messages=[{"role": "user", "content": prompt_text}], think=False, options={"temperature": 0.3, "num_predict": 80}) | |
| return resp["message"]["content"].strip() | |
| except Exception as e: return f'Ollama error: {e}' | |
| def process_image(img_path): | |
| if img_path is None: return 'No image received.', '[]', 'No image' | |
| try: | |
| pil_img = PILImage.open(img_path).convert('RGB') | |
| img_np = np.array(pil_img) | |
| except Exception as e: return f'Could not load image: {e}', '[]', 'Load error' | |
| t0 = time.time() | |
| dets, fallback = _detect(img_np) | |
| if fallback: | |
| _state['scene_json'] = '[]' | |
| return fallback, '[]', 'No objects detected' | |
| scene_str = json.dumps(dets, indent=2) | |
| _state['scene_json'] = scene_str | |
| rag_ctx = _retrieve(' '.join(d['label'] for d in dets[:4])) | |
| _state['rag_ctx'] = rag_ctx | |
| fact_lines = [] | |
| where_map = {'left':'to your left', 'center':'in front of you', 'right':'to your right'} | |
| for d in dets: | |
| color_part = f"{d['color']} " if d['color'] != 'unknown' else "" | |
| fact_lines.append(f"- a {color_part}{d['label']}, {where_map[d['position']]}") | |
| facts_str = '\n'.join(fact_lines) if fact_lines else '- Several objects detected.' | |
| prompt = (f"Act as a strictly factual descriptive assistant for a blind user.\nFacts:\n{facts_str}\nContext: {rag_ctx}\nRule 1: ONLY describe objects in the list above.\nRule 2: Write natural prose.\nRule 3: Start sentence 1 with the room type.\nRule 4: Output exactly 2 to 3 sentences total.") | |
| answer = _call_gemma_vision(prompt, pil_img) | |
| _state['last_answer'] = answer | |
| return answer, scene_str, f'✓ {len(dets)} objects | {round(time.time()-t0, 2)}s' | |
| def answer_question(q): | |
| if not q or not q.strip(): return 'Please type a question.' | |
| if _state['scene_json'] == '[]': return 'Please take a photo first.' | |
| dets = json.loads(_state['scene_json']) | |
| obj_str = ', '.join(f"{d.get('color','')} {d['label']} ({d['position']})" for d in dets) | |
| prompt = (f"A blind person is asking about the room.\nObjects detected: {obj_str}\nPrevious description: {_state['last_answer']}\nQuestion: {q}\nAnswer in 1 clear, natural sentence.") | |
| answer = _call_gemma_text(prompt) | |
| _state['last_answer'] = answer | |
| return answer | |
| with gr.Blocks(title='Assistive VQA — bonsAI') as demo: | |
| gr.Markdown('# 🦯 Assistive VQA (Local Deployment)') | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| img_in = gr.Image(label='Upload or capture photo', type='filepath', sources=['upload','webcam']) | |
| desc_btn = gr.Button('🔍 Describe Scene', variant='primary', size='lg') | |
| status = gr.Textbox(label='Status', interactive=False, value='Ready.') | |
| with gr.Column(scale=1): | |
| desc_out = gr.Textbox(label='Description', lines=6, interactive=False) | |
| with gr.Row(): | |
| q_in = gr.Textbox(label='Type your question', lines=1, scale=4) | |
| ask_btn = gr.Button('💬 Ask', variant='secondary', scale=1) | |
| ans_out = gr.Textbox(label='Answer', lines=2, interactive=False) | |
| with gr.Accordion('🔧 Debug — Scene JSON', open=False): | |
| json_out = gr.Code(language='json', lines=14) | |
| desc_btn.click(fn=process_image, inputs=[img_in], outputs=[desc_out, json_out, status]) | |
| ask_btn.click(fn=answer_question, inputs=[q_in], outputs=[ans_out]) | |
| q_in.submit(fn=answer_question, inputs=[q_in], outputs=[ans_out]) | |
| if __name__ == "__main__": | |
| print('Launching Gradio Server...') | |
| demo.launch(server_name="0.0.0.0", server_port=7860) |