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)