hshivhare/Visual_Aid_Designer
0
1import os2from selenium import webdriver3from selenium.webdriver.chrome.options import Options4from selenium.webdriver.common.by import By5from selenium.webdriver.support.ui import WebDriverWait6from selenium.webdriver.support import expected_conditions as EC7from tempfile import NamedTemporaryFile8from pathlib import Path9from selenium.webdriver.chrome.service import Service10from selenium.webdriver.chrome.options import Options11from selenium import webdriver12import chromedriver_autoinstaller13 14class MermaidRenderer:15 HTML_TEMPLATE = """16<!DOCTYPE html>17<html>18<head>19 <meta charset="UTF-8">20 <script src="https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.js"></script>21 <script>22 document.addEventListener("DOMContentLoaded", function() {{23 mermaid.initialize({{24 startOnLoad: true,25 theme: "default",26 themeVariables: {{27 primaryColor: "#ffffff",28 primaryTextColor: "#000000",29 primaryBorderColor: "#000000",30 lineColor: "#000000",31 secondaryColor: "#e0e0e0",32 tertiaryColor: "#f5f5f5",33 noteBkgColor: "#ffffff",34 noteTextColor: "#000000"35 }},36 flowchart: {{ curve: "basis", padding: 20 }}37 }});38 }});39 </script>40 <style>41 body {{ margin: 0; padding: 0; background: white; }}42 .mermaid {{ background: white; padding: 20px; display: inline-block; max-width: 4d00px; }}43 </style>44</head>45<body>46 <div class="mermaid">47 {code}48 </div>49</body>50</html>51"""52 53 54 55 # def setup_chrome_driver(self):56 # chromedriver_autoinstaller.install() # Automatically downloads and sets up ChromeDriver57 # chrome_options = Options()58 # chrome_options.add_argument("--headless")59 # chrome_options.add_argument("--no-sandbox")60 # chrome_options.add_argument("--disable-dev-shm-usage")61 # chrome_options.add_argument("--disable-gpu")62 # return webdriver.Chrome(options=chrome_options)63 64 def setup_chrome_driver(self):65 # Automatically install ChromeDriver66 chromedriver_autoinstaller.install()67 68 # Configure Chrome options69 chrome_options = Options()70 chrome_options.add_argument("--headless")71 chrome_options.add_argument("--no-sandbox")72 chrome_options.add_argument("--disable-dev-shm-usage")73 chrome_options.add_argument("--disable-gpu")74 75 # Set the path to Chromium explicitly76 chrome_options.binary_location = "/usr/bin/chromium-browser"77 78 # Return the WebDriver instance79 return webdriver.Chrome(options=chrome_options)80 81 82 def render_diagram(self, mermaid_code, output_path):83 driver = None84 temp_file_path = None85 86 try:87 with NamedTemporaryFile(delete=False, suffix=".html", mode="w", encoding="utf-8") as temp_file:88 temp_file.write(self.HTML_TEMPLATE.format(code=mermaid_code))89 temp_file_path = temp_file.name90 91 driver = self.setup_chrome_driver()92 driver.get(f"file://{temp_file_path}")93 94 wait = WebDriverWait(driver, 10)95 wait.until(EC.presence_of_element_located((By.CLASS_NAME, "mermaid")))96 97 width = driver.execute_script("return document.querySelector('.mermaid').offsetWidth")98 height = driver.execute_script("return document.querySelector('.mermaid').offsetHeight")99 driver.set_window_size(width + 40, height + 200)100 101 driver.save_screenshot(output_path)102 return True103 104 except Exception as e:105 print(f"Error: {e}")106 return False107 108 finally:109 if driver:110 driver.quit()111 if temp_file_path and os.path.exists(temp_file_path):112 os.unlink(temp_file_path)113 114 115if __name__ == "__main__":116 mermaid_code = """117 classDiagram118 class Bank {119 +String name120 +String address121 +List branches122 +addBranch(Branch branch)123 +removeBranch(Branch branch)124 }125 126 class Branch {127 +String branchId128 +String location129 +Bank bank130 +List employees131 +List accounts132 +addEmployee(Employee employee)133 +removeEmployee(Employee employee)134 +addAccount(Account account)135 +removeAccount(Account account)136 }137 138 class Employee {139 +String employeeId140 +String name141 +String position142 +Branch branch143 +assignToBranch(Branch branch)144 }145 146 class Customer {147 +String customerId148 +String name149 +String address150 +List accounts151 +addAccount(Account account)152 +removeAccount(Account account)153 }154 155 class Account {156 +String accountNumber157 +double balance158 +Customer owner159 +Transaction[] transactions160 +deposit(double amount)161 +withdraw(double amount)162 }163 164 class Transaction {165 +String transactionId166 +Date date167 +double amount168 +String type169 +Account account170 }171 172 class Loan {173 +String loanId174 +double amount175 +double interestRate176 +Branch branch177 +Customer borrower178 }179 180 Bank "1" o-- "many" Branch : has181 Branch "1" *-- "many" Account : contains182 Branch "1" o-- "many" Employee : employs183 Customer "1" *-- "many" Account : owns184 Account "1" -- "many" Transaction : contains185 Branch "1" *-- "many" Loan : processes186 Customer "1" *-- "many" Loan : borrows187 188 """189 190 renderer = MermaidRenderer()191 output_path = "mermaid_diagram.png"192 success = renderer.render_diagram(mermaid_code, output_path)193 194 if success:195 print(f"Diagram successfully generated at: {output_path}")196 else:197 print("Failed to generate diagram")198