Diagram2graph Adapters

🚀 Diagram-to-Graph Model

This model, a research-driven project developed during an internship at Zackariya Solution, specializes in extracting structured data (JSON) from images. It can specifically extract nodes, edges, and their sub-attributes to represent visual information as knowledge graphs.

⚠️ Important Note

This model is for learning purposes only and not for production applications. The extracted structured data may vary according to project needs.

🚀 Quick Start

%pip install -q "transformers>=4.49.0" accelerate datasets "qwen-vl-utils[decord]==0.0.8"

import os
import PIL
import torch
from qwen_vl_utils import process_vision_info
from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor


MODEL_ID="zackriya/diagram2graph-adapters"
MAX_PIXELS = 1280 * 28 * 28
MIN_PIXELS = 256 * 28 * 28


model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
	MODEL_ID,
	device_map="auto",
	torch_dtype=torch.bfloat16
)

processor = Qwen2_5_VLProcessor.from_pretrained(
	MODEL_ID,
	min_pixels=MIN_PIXELS,
	max_pixels=MAX_PIXELS
)


SYSTEM_MESSAGE = """You are a Vision Language Model specialized in extracting structured data from visual representations of process and flow diagrams.
Your task is to analyze the provided image of a diagram and extract the relevant information into a well-structured JSON format.
The diagram includes details such as nodes and edges. each of them have their own attributes.
Focus on identifying key data fields and ensuring the output adheres to the requested JSON structure.
Provide only the JSON output based on the extracted information. Avoid additional explanations or comments."""

def run_inference(image):
	messages= [
    	{
        	"role": "system",
        	"content": [{"type": "text", "text": SYSTEM_MESSAGE}],
    	},
    	{
        	"role": "user",
        	"content": [
            	{
                	"type": "image",
                	# this image is handled by qwen_vl_utils's process_visio_Info so no need to worry about pil image or path
                	"image": image,
            	},
            	{
                	"type": "text",
                	"text": "Extract data in JSON format, Only give the JSON",
            	},
        	],
    	},
	]

	text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
	image_inputs, _ = process_vision_info(messages)

	inputs = processor(
    	text=[text],
    	images=image_inputs,
    	return_tensors="pt",
	)
	inputs = inputs.to('cuda')

	generated_ids = model.generate(**inputs, max_new_tokens=512)
	generated_ids_trimmed = [
    	out_ids[len(in_ids):]
    	for in_ids, out_ids
    	in zip(inputs.input_ids, generated_ids)
	]

	output_text = processor.batch_decode(
    	generated_ids_trimmed,
    	skip_special_tokens=True,
    	clean_up_tokenization_spaces=False
	)
	return output_text
image = eval_dataset[9]['image'] # PIL image
# `image` could be URL or relative path to the image
output = run_inference(image)

# JSON loading
import json
json.loads(output[0])

✨ Features

✅ Direct Use

• Experiment with diagram-to-graph conversion 📊
• Understand AI-driven structured extraction from images

🚀 Downstream Use (Potential)

• Enhance BPMN/Flowchart analysis 🏗️
• Support automated document processing 📄

❌ Out-of-Scope Use

• Not designed for real-world production deployment ⚠️
• May not generalize well across all diagram types

📚 Documentation

📝 Model Details

🏗️ Training Details

• Dataset: Internally curated diagram dataset 🖼️
• Fine-tuning: LoRA-based optimization ⚡
• Precision: bf16 mixed-precision training 🎯

📈 Evaluation

• Metrics: F1-score 🏆
• Limitations: May struggle with complex, dense diagrams ⚠️

Results

• +14% improvement in node detection
• +23% improvement in edge detection

🤝 Collaboration

Are you interested in fine tuning your own model for your use case or want to explore how we can help you? Let's collaborate.

Zackriya Solutions

🔗 References

• Roboflow
• Qwen

🚀Stay Curious & Keep Exploring!🚀

📄 License

This project is licensed under the Apache 2.0 license.