graphvision-ai 0.1.0__tar.gz

graphvision_ai-0.1.0/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: graphvision-ai
+Version: 0.1.0
+Summary: Automatic Graph Classification and Data Extraction
+Author: Aryan Gahlot
+Description-Content-Type: text/markdown
+Requires-Dist: torch
+Requires-Dist: torchvision
+Requires-Dist: opencv-python
+Requires-Dist: easyocr
+Requires-Dist: Pillow
+Requires-Dist: numpy

graphvision_ai-0.1.0/graphvision/analyzer.py

@@ -0,0 +1,212 @@
+import torch
+import cv2
+import easyocr
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+
+from .model import (
+    load_classifier,
+    load_hbar_model,
+    load_vbar_model,
+    load_pie_model
+)
+
+class GraphAnalyzer:
+    def __init__(self):
+        # Automatically use Mac MPS, Nvidia CUDA, or CPU
+        self.device = torch.device("mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu"))
+
+        # 1. Load Models 
+        self.classifier = load_classifier(self.device)
+        self.hbar_model = load_hbar_model(self.device)
+        self.vbar_model = load_vbar_model(self.device)
+        self.pie_model = load_pie_model(self.device)
+
+        # 2. Setup OCR Engine
+        self.reader = easyocr.Reader(['en'])
+
+        # 3. CRITICAL: Setup strict ResNet Normalization
+        self.data_transforms = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+
+        # Adjust these indices to match exactly how the notebook classifier sorted the folders
+        self.class_map = {
+            0: "dot_line",
+            1: "hbar_categorical",
+            2: "line",
+            3: "pie",
+            4: "vbar_categorical"
+        }
+
+    def analyze(self, image):
+        """Main routing function. Classifies the image and extracts data."""
+        
+        # --- ROBUST IMAGE LOADING ---
+        if isinstance(image, str):
+            img_pil = Image.open(image).convert('RGB')
+            cv_img = cv2.imread(image)
+        elif isinstance(image, np.ndarray):
+            cv_img = image
+            img_pil = Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
+        elif isinstance(image, Image.Image):
+            img_pil = image.convert('RGB')
+            cv_img = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
+        else:
+            raise TypeError("Input must be a file path (str), NumPy array (cv2), or PIL Image.")
+
+        # 1. Classify the graph type
+        image_tensor = self.data_transforms(img_pil).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            outputs = self.classifier(image_tensor)
+            predicted_class = torch.argmax(outputs, dim=1).item()
+
+        graph_type = self.class_map.get(predicted_class, "unknown")
+
+        # 2. Route to the correct OCR + Prediction logic
+        if graph_type == "pie":
+            return self._extract_pie_data(cv_img, image_tensor)
+        elif graph_type == "vbar_categorical":
+            return self._extract_vbar_data(cv_img, image_tensor)
+        elif graph_type == "hbar_categorical":
+            return self._extract_hbar_data(cv_img, image_tensor)
+        else:
+            return {"type": graph_type, "status": f"Extraction logic for {graph_type} is not completely implemented in this snippet"}
+
+    def _extract_pie_data(self, cv_img, image_tensor):
+        """The dedicated Pie Chart extraction logic"""
+        with torch.no_grad():
+            preds = self.pie_model(image_tensor).squeeze().cpu().numpy() * 100.0
+            
+        h, w, _ = cv_img.shape
+        gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
+        all_text_results = self.reader.readtext(gray, mag_ratio=2.5)
+        
+        title = "Untitled"
+        raw_legend_names = [] 
+        
+        for bbox, text, conf in all_text_results:
+            clean_text = text.strip()
+            if not clean_text: continue
+                
+            x_center = (bbox[0][0] + bbox[2][0]) / 2
+            y_center = (bbox[0][1] + bbox[2][1]) / 2
+            x_pct = x_center / w
+            y_pct = y_center / h
+            
+            if y_pct < 0.15:
+                title = clean_text
+            elif y_pct > 0.15:
+                if len(clean_text) > 2 and not clean_text.replace('.', '', 1).isdigit():
+                    if clean_text.lower() == "grav": clean_text = "Gray"
+                    raw_legend_names.append((y_pct, clean_text)) 
+
+        legend_names = [item[1] for item in sorted(raw_legend_names, key=lambda i: i[0])]
+        
+        num_slices = len(legend_names)
+        if num_slices == 0:
+            valid_preds = [v for v in preds if v > 1.5]
+            num_slices = len(valid_preds)
+            
+        num_slices = min(num_slices, 10)
+        slice_preds = preds[:num_slices]
+        total_pred = sum(slice_preds)
+        
+        if total_pred > 0:
+            normalized_preds = [(v / total_pred) * 100.0 for v in slice_preds]
+        else:
+            normalized_preds = slice_preds
+        
+        final_slices = {}
+        for i in range(num_slices):
+            slice_name = legend_names[i] if i < len(legend_names) else f"Unknown_Slice_{i+1}"
+            slice_value = round(normalized_preds[i], 2) if i < len(normalized_preds) else 0.0
+            final_slices[slice_name] = slice_value
+
+        return {
+            "type": "pie",
+            "title": title,
+            "data": final_slices
+        }
+
+    def _extract_vbar_data(self, cv_img, image_tensor):
+        """The dedicated Vertical Bar Chart extraction logic"""
+        with torch.no_grad():
+            preds = self.vbar_model(image_tensor).squeeze().cpu().numpy() * 100.0
+            
+        h, w, _ = cv_img.shape
+        
+        # 1. Title
+        title_res = self.reader.readtext(cv_img[0:int(h*0.15), :])
+        title = title_res[0][1] if title_res else "Untitled"
+        
+        # 2. X-Axis Labels (Rotated crop to catch vertical text)
+        label_area = cv_img[int(h*0.70):h, :] 
+        label_rot = cv2.rotate(label_area, cv2.ROTATE_90_COUNTERCLOCKWISE)
+        label_res = self.reader.readtext(label_rot)
+        
+        x_labels = [res[1] for res in label_res if len(res[1]) > 1]
+        x_labels = x_labels[::-1]
+
+        # 3. Intelligent Matching (The dynamic fix)
+        max_val = np.max(preds)
+        noise_threshold = max_val * 0.05 
+        
+        valid_preds = [v for v in preds if v > noise_threshold]
+        
+        count = max(len(x_labels), len(valid_preds))
+        count = min(count, 10) 
+        
+        final_values = [round(float(v), 2) for v in preds[:count]]
+        
+        while len(x_labels) < len(final_values):
+            x_labels.append(f"Unknown_Label_{len(x_labels)+1}")
+
+        return {
+            "type": "vbar_categorical",
+            "title": title,
+            "x_axis_labels": x_labels,
+            "values": final_values
+        }
+
+    def _extract_hbar_data(self, cv_img, image_tensor):
+        """The dedicated Horizontal Bar Chart extraction logic"""
+        with torch.no_grad():
+            preds = self.hbar_model(image_tensor).squeeze().cpu().numpy() * 100.0
+            
+        h, w, _ = cv_img.shape
+        
+        # 1. Title
+        title_res = self.reader.readtext(cv_img[0:int(h*0.15), :])
+        title = title_res[0][1] if title_res else "Untitled"
+        
+        # 2. Labels (Smooth Contrast applied)
+        label_area = cv_img[int(h*0.10):int(h*0.95), 0:int(w*0.45)] 
+        gray = cv2.cvtColor(label_area, cv2.COLOR_BGR2GRAY)
+        contrast_img = cv2.convertScaleAbs(gray, alpha=1.5, beta=-50)
+        
+        label_res = self.reader.readtext(contrast_img, mag_ratio=2.0)
+        y_labels = [res[1] for res in label_res if len(res[1]) > 1]
+        y_labels = y_labels[::-1] 
+
+        # 3. Intelligent Matching
+        valid_preds = [v for v in preds if v > 2.0]
+        
+        count = len(y_labels) if len(y_labels) > 0 else len(valid_preds)
+        count = min(count, 10) 
+        
+        final_values = [round(float(v), 2) for v in preds[:count]]
+        
+        while len(y_labels) < len(final_values):
+            y_labels.append(f"Unknown_Label_{len(y_labels)+1}")
+
+        return {
+            "type": "hbar_categorical",
+            "title": title,
+            "y_axis_labels": y_labels,
+            "values": final_values
+        }

graphvision_ai-0.1.0/graphvision/model.py

@@ -0,0 +1,90 @@
+import torch
+import torch.nn as nn
+from torchvision import models
+import os
+import urllib.request
+import shutil
+
+# --- AUTO-DOWNLOAD WEIGHTS SETUP ---
+HOME_DIR = os.path.expanduser("~")
+# Create a hidden folder in the user's home directory to store weights securely
+WEIGHTS_DIR = os.path.join(HOME_DIR, ".graphvision_weights")
+os.makedirs(WEIGHTS_DIR, exist_ok=True)
+
+# Your Hugging Face Repository
+HF_BASE_URL = "https://huggingface.co/ShadowGard3n/graphvision/resolve/main/"
+
+WEIGHT_URLS = {
+    "graph_classifier.pth": f"{HF_BASE_URL}graph_classifier.pth",
+    "pie_regressor_stable.pth": f"{HF_BASE_URL}pie_regressor_stable.pth",
+    "hbar_regressor_stable.pth": f"{HF_BASE_URL}hbar_regressor_stable.pth",
+    "vbar_regressor.pth": f"{HF_BASE_URL}vbar_regressor.pth"
+}
+
+def download_weight(filename):
+    """Downloads the weight file from Hugging Face if it doesn't already exist."""
+    filepath = os.path.join(WEIGHTS_DIR, filename)
+    if not os.path.exists(filepath):
+        print(f"Downloading {filename} from Hugging Face (This only happens once)...")
+        # Add a custom User-Agent header so Hugging Face doesn't block the request
+        req = urllib.request.Request(WEIGHT_URLS[filename], headers={'User-Agent': 'Mozilla/5.0'})
+        with urllib.request.urlopen(req) as response, open(filepath, 'wb') as out_file:
+            shutil.copyfileobj(response, out_file)
+    return filepath
+
+# --- 1. MODEL ARCHITECTURES ---
+class PieRegressor(nn.Module):
+    def __init__(self):
+        super(PieRegressor, self).__init__()
+        self.backbone = models.resnet18(weights=None)
+        num_ftrs = self.backbone.fc.in_features
+        self.backbone.fc = nn.Linear(num_ftrs, 10)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        return self.sigmoid(self.backbone(x))
+
+class BarRegressor(nn.Module):
+    def __init__(self):
+        super(BarRegressor, self).__init__()
+        self.backbone = models.resnet18(weights=None)
+        num_ftrs = self.backbone.fc.in_features
+        self.backbone.fc = nn.Linear(num_ftrs, 10)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        return self.sigmoid(self.backbone(x))
+
+# --- 2. SECURE LOADING FUNCTIONS ---
+def load_classifier(device):
+    # FIX: Use the download_weight function instead of os.path.join
+    path = download_weight("graph_classifier.pth")
+    
+    model = models.resnet18(weights=None)
+    num_ftrs = model.fc.in_features
+    # Set to 5 classes to perfectly match the saved checkpoint
+    model.fc = nn.Linear(num_ftrs, 5) 
+    
+    model.load_state_dict(torch.load(path, map_location=device))
+    return model.to(device).eval()
+
+def load_pie_model(device):
+    # FIX: Use the download_weight function
+    path = download_weight("pie_regressor_stable.pth")
+    model = PieRegressor()
+    model.load_state_dict(torch.load(path, map_location=device))
+    return model.to(device).eval()
+
+def load_hbar_model(device):
+    # FIX: Use the download_weight function
+    path = download_weight("hbar_regressor_stable.pth")
+    model = BarRegressor()
+    model.load_state_dict(torch.load(path, map_location=device))
+    return model.to(device).eval()
+
+def load_vbar_model(device):
+    # FIX: Use the download_weight function
+    path = download_weight("vbar_regressor.pth") 
+    model = BarRegressor()
+    model.load_state_dict(torch.load(path, map_location=device))
+    return model.to(device).eval()

graphvision_ai-0.1.0/graphvision_ai.egg-info/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: graphvision-ai
+Version: 0.1.0
+Summary: Automatic Graph Classification and Data Extraction
+Author: Aryan Gahlot
+Description-Content-Type: text/markdown
+Requires-Dist: torch
+Requires-Dist: torchvision
+Requires-Dist: opencv-python
+Requires-Dist: easyocr
+Requires-Dist: Pillow
+Requires-Dist: numpy

graphvision_ai-0.1.0/graphvision_ai.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+README.md
+pyproject.toml
+graphvision/__init__.py
+graphvision/analyzer.py
+graphvision/model.py
+graphvision_ai.egg-info/PKG-INFO
+graphvision_ai.egg-info/SOURCES.txt
+graphvision_ai.egg-info/dependency_links.txt
+graphvision_ai.egg-info/requires.txt
+graphvision_ai.egg-info/top_level.txt

graphvision_ai-0.1.0/graphvision_ai.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

graphvision_ai-0.1.0/graphvision_ai.egg-info/requires.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+opencv-python
+easyocr
+Pillow
+numpy

graphvision_ai-0.1.0/graphvision_ai.egg-info/top_level.txt

@@ -0,0 +1 @@
+graphvision

graphvision_ai-0.1.0/pyproject.toml

@@ -0,0 +1,21 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "graphvision-ai"  
+version = "0.1.0"
+description = "Automatic Graph Classification and Data Extraction"
+readme = "README.md"
+authors = [{name="Aryan Gahlot"}]
+dependencies = [
+    "torch",
+    "torchvision",
+    "opencv-python",
+    "easyocr",
+    "Pillow",
+    "numpy"
+]
+
+[tool.setuptools.packages.find]
+include = ["graphvision"]

graphvision_ai-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+