sightrag 0.1.0__py3-none-any.whl

sightrag/__init__.py

@@ -0,0 +1,31 @@
+"""
+SightRAG — Image and Video RAG
+See. Search. Retrieve.
+
+pip install sightrag
+
+Usage:
+    from sightrag import SightRAG
+
+    rag = SightRAG()
+    rag.index("./photos/")
+    results = rag.query("find empty shelf")
+
+REST API:
+    from sightrag import serve
+    serve(port=8000)
+"""
+
+from .core import SightRAG
+
+__version__ = "0.1.0"
+__author__ = "Ant (VK-Ant)"
+
+
+def serve(host: str = "0.0.0.0", port: int = 8000):
+    """Start SightRAG REST API server."""
+    from .api import serve as _serve
+    _serve(host=host, port=port)
+
+
+__all__ = ["SightRAG", "serve"]

sightrag/api.py

@@ -0,0 +1,189 @@
+# sightrag/api.py
+# REST API — FastAPI based
+# Run: sightrag-server or python -m sightrag.api
+
+import os
+import json
+import shutil
+import tempfile
+from pathlib import Path
+
+try:
+    from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+    from fastapi.responses import JSONResponse
+    import uvicorn
+except ImportError:
+    raise ImportError(
+        "FastAPI not installed.\n"
+        "Run: pip install sightrag[api]"
+    )
+
+from .core import SightRAG
+
+app = FastAPI(
+    title="SightRAG API",
+    description="See. Search. Retrieve. — Image and Video RAG",
+    version="0.1.0"
+)
+
+# Global SightRAG instance
+rag = None
+
+
+def get_rag():
+    global rag
+    if rag is None:
+        store = os.getenv("SIGHTRAG_STORE", "sqlite")
+        domain = os.getenv("SIGHTRAG_DOMAIN", None)
+        index_path = os.getenv("SIGHTRAG_INDEX", "./sightrag_index")
+        rag = SightRAG(store=store, domain_hint=domain, index_path=index_path)
+    return rag
+
+
+@app.get("/")
+def root():
+    return {
+        "name": "SightRAG API",
+        "version": "0.1.0",
+        "tagline": "See. Search. Retrieve.",
+        "endpoints": {
+            "POST /index/folder": "Index an image folder",
+            "POST /index/video": "Index a video file",
+            "POST /index/upload": "Upload and index images",
+            "POST /query/text": "Search with text",
+            "POST /query/reference": "Search with reference image",
+            "GET /status": "Index status",
+            "DELETE /index": "Clear index"
+        }
+    }
+
+
+@app.get("/status")
+def status():
+    r = get_rag()
+    return {
+        "indexed_regions": r.count(),
+        "store": r._store_type,
+        "domain_hint": r.domain_hint
+    }
+
+
+@app.post("/index/folder")
+def index_folder(path: str = Form(...)):
+    """Index all images in a folder."""
+    r = get_rag()
+    try:
+        r.index(path)
+        return {
+            "status": "success",
+            "indexed_regions": r.count(),
+            "source": path
+        }
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
+@app.post("/index/video")
+def index_video(path: str = Form(...), fps: int = Form(1)):
+    """Index a video file."""
+    r = get_rag()
+    try:
+        r.index(path, fps=fps)
+        return {
+            "status": "success",
+            "indexed_regions": r.count(),
+            "source": path,
+            "fps": fps
+        }
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
+@app.post("/index/upload")
+async def index_upload(files: list[UploadFile] = File(...)):
+    """Upload and index images directly."""
+    r = get_rag()
+    upload_dir = tempfile.mkdtemp(prefix="sightrag_upload_")
+
+    try:
+        # Save uploaded files
+        for f in files:
+            file_path = os.path.join(upload_dir, f.filename)
+            with open(file_path, "wb") as out:
+                content = await f.read()
+                out.write(content)
+
+        # Index the upload folder
+        r.index(upload_dir)
+
+        return {
+            "status": "success",
+            "files_uploaded": len(files),
+            "indexed_regions": r.count()
+        }
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    finally:
+        shutil.rmtree(upload_dir, ignore_errors=True)
+
+
+@app.post("/query/text")
+def query_text(text: str = Form(...), top_k: int = Form(5)):
+    """Search with plain English text."""
+    r = get_rag()
+    try:
+        results = r.query(text=text, top_k=top_k)
+        return {
+            "query": text,
+            "results": results,
+            "count": len(results)
+        }
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
+@app.post("/query/reference")
+async def query_reference(
+    file: UploadFile = File(...),
+    top_k: int = Form(5)
+):
+    """Search using a reference image."""
+    r = get_rag()
+
+    # Save reference temporarily
+    tmp = tempfile.NamedTemporaryFile(
+        delete=False, suffix=f"_{file.filename}"
+    )
+    try:
+        content = await file.read()
+        tmp.write(content)
+        tmp.close()
+
+        results = r.query(reference=tmp.name, top_k=top_k)
+        return {
+            "reference": file.filename,
+            "results": results,
+            "count": len(results)
+        }
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    finally:
+        os.unlink(tmp.name)
+
+
+@app.delete("/index")
+def clear_index():
+    """Clear all indexed data."""
+    r = get_rag()
+    r.clear()
+    return {"status": "cleared", "indexed_regions": 0}
+
+
+def serve(host: str = "0.0.0.0", port: int = 8000):
+    """Start the SightRAG API server."""
+    print(f"[SightRAG] Starting API server on {host}:{port}")
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    serve()

sightrag/core.py

@@ -0,0 +1,129 @@
+# sightrag/core.py
+# Main SightRAG class
+# All data stored in ~/.sightrag/ — project folder stays clean
+
+import os
+from pathlib import Path
+from .detector import Detector
+from .embedder import Embedder
+from .indexer import Indexer
+from .retriever import Retriever
+
+SIGHTRAG_HOME = os.path.join(Path.home(), ".sightrag")
+
+
+class SightRAG:
+    """
+    SightRAG — Image and Video RAG.
+    See. Search. Retrieve.
+
+    Usage:
+        rag = SightRAG()
+        rag.index("./photos/")
+        results = rag.query("find empty shelf")
+    """
+
+    def __init__(self,
+                 store: str = "sqlite",
+                 domain_hint: str = None,
+                 index_path: str = None):
+
+        self.domain_hint = domain_hint
+        self._store_type = store
+
+        if index_path is None:
+            self._index_path = os.path.join(SIGHTRAG_HOME, "index")
+        else:
+            self._index_path = index_path
+
+        os.makedirs(SIGHTRAG_HOME, exist_ok=True)
+
+        print("[SightRAG] Initializing...")
+        self._detector = Detector()
+        self._embedder = Embedder()
+        self._store = self._init_store(store, self._index_path)
+        self._indexer = Indexer(
+            self._detector, self._embedder,
+            self._store, domain_hint
+        )
+        self._retriever = Retriever(
+            self._embedder, self._detector,
+            self._store, domain_hint
+        )
+        print("[SightRAG] Ready.")
+
+    def _init_store(self, store_type: str, path: str):
+        if store_type == "chroma":
+            try:
+                from .store.chroma_store import ChromaStore
+                return ChromaStore(path)
+            except ImportError:
+                print("[SightRAG] ChromaDB not found. Using SQLite.")
+                from .store.sqlite_store import SQLiteStore
+                return SQLiteStore(path)
+        elif store_type == "sqlite":
+            from .store.sqlite_store import SQLiteStore
+            return SQLiteStore(path)
+        else:
+            raise ValueError(f"Unknown store: {store_type}. Use 'chroma' or 'sqlite'.")
+
+    def index(self, path: str = None, source: str = None,
+              camera_id: int = 0, fps: int = 1):
+        """Index images, video, or camera."""
+        if source == "camera":
+            self._indexer.index_camera(camera_id=camera_id, fps=fps)
+            return self
+
+        if path is None:
+            raise ValueError("Provide a path or source='camera'")
+
+        if os.path.isdir(path):
+            self._indexer.index_folder(path, fps=fps)
+        elif os.path.isfile(path):
+            ext = os.path.splitext(path)[1].lower()
+            if ext in {".mp4", ".avi", ".mov", ".mkv"}:
+                self._indexer.index_video(path, fps=fps)
+            else:
+                from .utils.image import load_image
+                image = load_image(path)
+                self._index_single_image(path, image)
+                print(f"[SightRAG] 1 image indexed. Total: {self.count()} regions.")
+        else:
+            raise FileNotFoundError(f"Path not found: {path}")
+
+        return self
+
+    def _index_single_image(self, path, image):
+        """Index one image with detection + embedding."""
+        import numpy as np
+        regions = self._detector.detect(image)
+        for j, region in enumerate(regions):
+            embedding = self._embedder.embed_image(region["crop"])
+            if not np.allclose(embedding, 0):
+                self._store.add(f"img_{j}", embedding, {
+                    "image_path":  str(path),
+                    "bbox":        region["bbox"],
+                    "label":       region["label"],
+                    "confidence":  region["confidence"],
+                    "source_type": "image"
+                })
+
+    def query(self, text: str = None, reference: str = None, top_k: int = 5):
+        """Search indexed content with text or reference image."""
+        if text is None and reference is None:
+            raise ValueError("Provide text or reference image.")
+        if text:
+            return self._retriever.query_text(text, top_k)
+        else:
+            return self._retriever.query_reference(reference, top_k)
+
+    def count(self) -> int:
+        return self._store.count()
+
+    def clear(self):
+        self._store.clear()
+        print("[SightRAG] Index cleared.")
+        return self
+
+    def __repr__(self):
+        return f"SightRAG(store='{self._store_type}', indexed={self.count()} regions)"

sightrag/detector.py

@@ -0,0 +1,81 @@
+# sightrag/detector.py
+# YOLO detection — models stored in ~/.sightrag/models/
+
+import os
+import numpy as np
+from PIL import Image
+from pathlib import Path
+
+
+MODEL_DIR = os.path.join(Path.home(), ".sightrag", "models")
+
+
+class Detector:
+    """YOLO object detector with whole-image fallback."""
+
+    def __init__(self, model_size: str = "yolo11n.pt", model_dir: str = None):
+        self.model = None
+        self.model_size = model_size
+        self.model_dir = model_dir or MODEL_DIR
+        os.makedirs(self.model_dir, exist_ok=True)
+        self._load()
+
+    def _load(self):
+        try:
+            from ultralytics import YOLO
+            import logging
+            logging.getLogger("ultralytics").setLevel(logging.WARNING)
+
+            model_path = os.path.join(self.model_dir, self.model_size)
+
+            if os.path.exists(model_path):
+                self.model = YOLO(model_path)
+            else:
+                # Download and move to our folder
+                self.model = YOLO(self.model_size)
+                # Move .pt file from current dir to model_dir
+                cwd_model = os.path.join(os.getcwd(), self.model_size)
+                if os.path.exists(cwd_model) and cwd_model != model_path:
+                    import shutil
+                    shutil.move(cwd_model, model_path)
+
+        except Exception as e:
+            print(f"[SightRAG] YOLO not available: {str(e)[:100]}")
+            self.model = None
+
+    def detect(self, image: Image.Image, confidence: float = 0.25):
+        """Detect objects. Always returns at least whole image."""
+        regions = []
+
+        if self.model is not None:
+            try:
+                results = self.model(image, conf=confidence, verbose=False)
+                for result in results:
+                    if result.boxes is None or len(result.boxes) == 0:
+                        continue
+                    for box in result.boxes:
+                        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+                        w, h = image.size
+                        x1, y1 = max(0, x1), max(0, y1)
+                        x2, y2 = min(w, x2), min(h, y2)
+                        if (x2 - x1) < 10 or (y2 - y1) < 10:
+                            continue
+                        regions.append({
+                            "crop":       image.crop((x1, y1, x2, y2)),
+                            "bbox":       [x1, y1, x2, y2],
+                            "label":      result.names[int(box.cls[0])],
+                            "confidence": float(box.conf[0])
+                        })
+            except:
+                pass
+
+        # Always add whole image
+        w, h = image.size
+        regions.append({
+            "crop":       image,
+            "bbox":       [0, 0, w, h],
+            "label":      "whole_image",
+            "confidence": 1.0
+        })
+
+        return regions

sightrag/embedder.py

@@ -0,0 +1,83 @@
+# sightrag/embedder.py
+# CLIP embedder — works across all transformers versions
+
+import os
+import numpy as np
+from PIL import Image
+
+os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+class Embedder:
+    """CLIP embedder for images and text queries."""
+
+    def __init__(self, model_name: str = "openai/clip-vit-base-patch32"):
+        self.model = None
+        self.processor = None
+        self.embed_dim = None
+        self._load(model_name)
+
+    def _load(self, model_name):
+        import warnings
+        warnings.filterwarnings("ignore")
+
+        from transformers import CLIPModel, CLIPProcessor
+        self.model = CLIPModel.from_pretrained(model_name)
+        self.processor = CLIPProcessor.from_pretrained(model_name)
+        self.model.eval()
+        self.embed_dim = self.model.config.projection_dim
+
+    def embed_image(self, image: Image.Image) -> np.ndarray:
+        """Embed image → fixed-size normalized vector."""
+        import torch
+        try:
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+
+            # Get pixel values
+            inputs = self.processor(images=image, return_tensors="pt", padding=True)
+            pixel_values = inputs["pixel_values"]
+
+            with torch.no_grad():
+                # Use vision model + projection explicitly
+                # This guarantees correct output dim across all versions
+                vision_out = self.model.vision_model(pixel_values=pixel_values)
+                pooled = vision_out.pooler_output            # (1, hidden_dim)
+                projected = self.model.visual_projection(pooled)  # (1, projection_dim)
+
+            emb = projected[0].detach().cpu().numpy().astype(np.float32)
+            norm = np.linalg.norm(emb)
+            return emb / norm if norm > 0 else emb
+
+        except Exception as e:
+            print(f"[SightRAG] Image embed error: {e}")
+            return np.zeros(self.embed_dim, dtype=np.float32)
+
+    def embed_text(self, text: str, domain_hint: str = None) -> np.ndarray:
+        """Embed text query → fixed-size normalized vector."""
+        import torch
+        try:
+            query = f"{text} {domain_hint}" if domain_hint else text
+
+            inputs = self.processor(
+                text=[query], return_tensors="pt",
+                padding=True, truncation=True
+            )
+
+            with torch.no_grad():
+                # Use text model + projection explicitly
+                text_out = self.model.text_model(
+                    input_ids=inputs["input_ids"],
+                    attention_mask=inputs["attention_mask"]
+                )
+                pooled = text_out.pooler_output              # (1, hidden_dim)
+                projected = self.model.text_projection(pooled)    # (1, projection_dim)
+
+            emb = projected[0].detach().cpu().numpy().astype(np.float32)
+            norm = np.linalg.norm(emb)
+            return emb / norm if norm > 0 else emb
+
+        except Exception as e:
+            print(f"[SightRAG] Text embed error: {e}")
+            return np.zeros(self.embed_dim, dtype=np.float32)

sightrag/indexer.py

@@ -0,0 +1,147 @@
+# sightrag/indexer.py
+# Image, video, camera indexing — clean output only
+
+import os
+import numpy as np
+from pathlib import Path
+from .detector import Detector
+from .embedder import Embedder
+from .utils.image import load_image, SUPPORTED_FORMATS as IMAGE_FORMATS
+from .utils.video import extract_frames, SUPPORTED_FORMATS as VIDEO_FORMATS
+
+
+class Indexer:
+    """Indexes images, videos, and camera frames."""
+
+    def __init__(self, detector, embedder, store, domain_hint=None):
+        self.detector = detector
+        self.embedder = embedder
+        self.store = store
+        self.domain_hint = domain_hint
+
+    def _index_image(self, path_str, image, prefix):
+        """Index one image — detect regions, embed, store."""
+        regions = self.detector.detect(image)
+        count = 0
+        for j, region in enumerate(regions):
+            embedding = self.embedder.embed_image(region["crop"])
+            if not np.allclose(embedding, 0):
+                self.store.add(f"{prefix}_{j}", embedding, {
+                    "image_path":  path_str,
+                    "bbox":        region["bbox"],
+                    "label":       region["label"],
+                    "confidence":  region["confidence"],
+                    "source_type": "image"
+                })
+                count += 1
+        return count
+
+    def index_folder(self, folder_path: str, fps: int = 1):
+        """Index all images AND videos in a folder."""
+        folder = Path(folder_path)
+        if not folder.exists():
+            raise FileNotFoundError(f"Folder not found: {folder}")
+        if not folder.is_dir():
+            raise ValueError(f"Not a folder: {folder}")
+
+        # Find images
+        image_paths = []
+        for fmt in IMAGE_FORMATS:
+            image_paths.extend(folder.glob(f"*{fmt}"))
+            image_paths.extend(folder.glob(f"*{fmt.upper()}"))
+        image_paths = sorted(set(image_paths))
+
+        # Find videos
+        video_paths = []
+        for fmt in VIDEO_FORMATS:
+            video_paths.extend(folder.glob(f"*{fmt}"))
+            video_paths.extend(folder.glob(f"*{fmt.upper()}"))
+        video_paths = sorted(set(video_paths))
+
+        if not image_paths and not video_paths:
+            raise ValueError(f"No images or videos in {folder}")
+
+        print(f"[SightRAG] Found {len(image_paths)} images, {len(video_paths)} videos")
+
+        # Index images
+        if image_paths:
+            total = len(image_paths)
+            for i, path in enumerate(image_paths, 1):
+                try:
+                    image = load_image(str(path))
+                    self._index_image(str(path), image, path.stem)
+                    pct = int((i / total) * 40)
+                    bar = "█" * pct + "░" * (40 - pct)
+                    print(f"\r  [{bar}] {i}/{total} images", end="", flush=True)
+                except Exception as e:
+                    print(f"\n  Skipping {path.name}: {e}")
+            print()
+
+        # Index videos
+        if video_paths:
+            for v_idx, vpath in enumerate(video_paths, 1):
+                try:
+                    print(f"[SightRAG] Video {v_idx}/{len(video_paths)}: {vpath.name}")
+                    self._index_video(str(vpath), fps)
+                except Exception as e:
+                    print(f"  Skipping {vpath.name}: {e}")
+
+        print(f"[SightRAG] Done. {self.store.count()} regions indexed.")
+
+    def index_video(self, video_path: str, fps: int = 1):
+        """Index a single video file."""
+        self._index_video(video_path, fps)
+        print(f"[SightRAG] Done. {self.store.count()} regions indexed.")
+
+    def _index_video(self, video_path: str, fps: int = 1):
+        """Internal video indexing."""
+        video_name = Path(video_path).stem
+        frames = extract_frames(video_path, fps=fps)
+        total = len(frames)
+        print(f"  {total} frames extracted...")
+
+        for i, (image, timestamp) in enumerate(frames, 1):
+            try:
+                regions = self.detector.detect(image)
+                for j, region in enumerate(regions):
+                    embedding = self.embedder.embed_image(region["crop"])
+                    if not np.allclose(embedding, 0):
+                        self.store.add(f"{video_name}_f{i}_r{j}", embedding, {
+                            "image_path":  video_path,
+                            "bbox":        region["bbox"],
+                            "label":       region["label"],
+                            "confidence":  region["confidence"],
+                            "timestamp":   timestamp,
+                            "source_type": "video"
+                        })
+                pct = int((i / total) * 40)
+                bar = "█" * pct + "░" * (40 - pct)
+                print(f"\r  [{bar}] {i}/{total} frames", end="", flush=True)
+            except:
+                pass
+        print()
+
+    def index_camera(self, camera_id=0, fps=1, buffer_seconds=60):
+        """Index live camera frames."""
+        from .utils.camera import capture_frames
+        print(f"[SightRAG] Camera {camera_id}. Press Ctrl+C to stop.")
+
+        count = 0
+        try:
+            for image, timestamp in capture_frames(camera_id, fps, buffer_seconds):
+                regions = self.detector.detect(image)
+                for j, region in enumerate(regions):
+                    embedding = self.embedder.embed_image(region["crop"])
+                    if not np.allclose(embedding, 0):
+                        self.store.add(f"cam{camera_id}_{timestamp}_{j}", embedding, {
+                            "image_path":  f"camera_{camera_id}",
+                            "bbox":        region["bbox"],
+                            "label":       region["label"],
+                            "confidence":  region["confidence"],
+                            "timestamp":   timestamp,
+                            "source_type": "camera"
+                        })
+                count += 1
+                print(f"\r[SightRAG] {count} frames | {timestamp}", end="", flush=True)
+        except KeyboardInterrupt:
+            print(f"\n[SightRAG] Stopped. {self.store.count()} regions indexed.")