simonduerr/pdb2vector
1
1 2import os3import json4import uuid5 6import gradio as gr7from gradio_moleculeview import moleculeview8 9 10import cellscape11 12 13def html_output(input_file):14 with open(input_file, "r") as f:15 svg = f.read().replace("<svg", "<svg id='svgElement'")16 17 x = (18 """<!DOCTYPE html>19<html lang="en">20 <head>21 <meta charset="UTF-8" />22 <meta name="viewport" content="width=device-width, initial-scale=1.0" />23 <title>Document</title>24 <script src="https://cdn.tailwindcss.com"></script>25 26 <style>27 #snackbar {28 visibility: hidden;29 color: #fff;30 background-color: #333;31 min-width: 250px;32 margin-left: -125px;33 border-radius: 2px;34 padding: 16px;35 text-align: center;36 left: 50%;37 bottom: 30px;38 z-index: 1;39 position: fixed;40}41 42/* This will be activated when the snackbar's class is 'show' which will be added through JS */43#snackbar.show {44 visibility: visible;45 -webkit-animation: fadein 0.5s, fadeout 0.5s 2.5s;46 animation: fadein 0.5s, fadeout 0.5s 2.5s;47}48 49/* Animations for fading in and out */50@-webkit-keyframes fadein {51 from {bottom: 0; opacity: 0;}52 to {bottom: 30px; opacity: 1;}53}54 55@keyframes fadein {56 from {bottom: 0; opacity: 0;}57 to {bottom: 30px; opacity: 1;}58}59 60@-webkit-keyframes fadeout {61 from {bottom: 30px; opacity: 1;}62 to {bottom: 0; opacity: 0;}63}64 65@keyframes fadeout {66 from {bottom: 30px; opacity: 1;}67 to {bottom: 0; opacity: 0;}68}69 </style>70 </head>71 <body>72 73 <div class="flex container mx-auto max-w-3xl">74 <div class="w-1/2">75 76 77 78 79 """80 + svg81 + """82 </div>83 84 <div class="w-1/2 space-y-4 justify-center items-center flex flex-col">85 <div class="w-3/4">86 <div87 class="w-full flex space-x-2 items-center text-xl justify-center font-light mb-3"88 >89 <svg90 xmlns="http://www.w3.org/2000/svg"91 class="w-4 h-4"92 viewBox="0 0 16 16"93 >94 <path95 fill="currentColor"96 fill-rule="evenodd"97 d="M6.091 5.194q.2-.322.305-.694h3.208c.259.916.98 1.637 1.896 1.896v3.208A2.76 2.76 0 0 0 9.604 11.5H6.396A2.76 2.76 0 0 0 4.5 9.604V6.396a2.76 2.76 0 0 0 1.591-1.202M13 6.396v3.208q.372.106.694.305a2.75 2.75 0 0 1 .897 3.785A2.751 2.751 0 0 1 9.603 13H6.397q-.106.372-.305.694A2.75 2.75 0 1 1 3 9.604V6.396A2.751 2.751 0 1 1 6.396 3h3.208A2.751 2.751 0 1 1 13 6.396M3.75 2.5a1.25 1.25 0 1 1 0 2.5a1.25 1.25 0 0 1 0-2.5M5 12.25a1.25 1.25 0 1 0-2.5 0a1.25 1.25 0 0 0 2.5 0M12.25 11a1.25 1.25 0 1 1 0 2.5a1.25 1.25 0 0 1 0-2.5m1.25-7.25a1.25 1.25 0 1 0-2.5 0a1.25 1.25 0 0 0 2.5 0"98 clip-rule="evenodd"99 />100 </svg>101 <div>SVG</div>102 </div>103 <div class="w-full flex space-x-5">104 <button105 id="copySvgBtn"106 class="flex items-center justify-center space-x-2 w-full p-3 border border-gray-300 rounded bg-gray-100"107 >108 <svg109 xmlns="http://www.w3.org/2000/svg"110 class="w-4 h-4"111 viewBox="0 0 16 16"112 >113 <path114 fill="currentColor"115 fill-rule="evenodd"116 d="M12 2.5H8A1.5 1.5 0 0 0 6.5 4v1H8a3 3 0 0 1 3 3v1.5h1A1.5 1.5 0 0 0 13.5 8V4A1.5 1.5 0 0 0 12 2.5M11 11h1a3 3 0 0 0 3-3V4a3 3 0 0 0-3-3H8a3 3 0 0 0-3 3v1H4a3 3 0 0 0-3 3v4a3 3 0 0 0 3 3h4a3 3 0 0 0 3-3zM4 6.5h4A1.5 1.5 0 0 1 9.5 8v4A1.5 1.5 0 0 1 8 13.5H4A1.5 1.5 0 0 1 2.5 12V8A1.5 1.5 0 0 1 4 6.5"117 clip-rule="evenodd"118 />119 </svg>120 <div>Copy*</div>121 </button>122 <button123 id="downloadSvgBtn"124 class="w-full flex items-center justify-center space-x-2 p-3 border border-gray-300 rounded bg-gray-100"125 >126 <svg127 xmlns="http://www.w3.org/2000/svg"128 class="h-4 w-4"129 viewBox="0 0 16 16"130 >131 <path132 fill="currentColor"133 fill-rule="evenodd"134 d="M6.5 6H4.029l3.499 4.276a.61.61 0 0 0 .944 0L11.971 6H9.5V2.75a.25.25 0 0 0-.25-.25h-2.5a.25.25 0 0 0-.25.25zm1.498 7.5H1.75a.75.75 0 0 0 0 1.5h12.5a.75.75 0 0 0 0-1.5zm0-1.5a2.11 2.11 0 0 1-1.631-.774L2.833 6.907A1.474 1.474 0 0 1 3.973 4.5H5V2.75C5 1.784 5.784 1 6.75 1h2.5c.966 0 1.75.784 1.75 1.75V4.5h1.026a1.474 1.474 0 0 1 1.14 2.407l-3.533 4.319c-.4.49-1 .774-1.632.774H8"135 clip-rule="evenodd"136 />137 </svg>138 <div>Download</div>139 </button>140 </div>141 142 <div class="text-xs mt-3 text-center w-full">143 *Only works in in Adobe Illustrator.144 </div>145 146 <div147 class="flex space-x-2 items-center text-xl justify-center font-light mb-3 mt-10"148 >149 <svg150 xmlns="http://www.w3.org/2000/svg"151 class="w-4 h-4"152 viewBox="0 0 16 16"153 >154 <path155 fill="currentColor"156 fill-rule="evenodd"157 d="M11.5 3h-7A1.5 1.5 0 0 0 3 4.5v5.027l.962-.7a1.75 1.75 0 0 1 2.079.016l.928.696l2.368-2.03a1.75 1.75 0 0 1 2.325.043L13 8.787V4.5A1.5 1.5 0 0 0 11.5 3m3 7.498V4.5a3 3 0 0 0-3-3h-7a3 3 0 0 0-3 3v7a3 3 0 0 0 3 3h7a3 3 0 0 0 3-3zm-1.5.33l-2.355-2.174a.25.25 0 0 0-.332-.006L7.488 11.07l-.457.392l-.481-.361l-1.41-1.057a.25.25 0 0 0-.296-.002L3 11.381v.119A1.5 1.5 0 0 0 4.5 13h7a1.5 1.5 0 0 0 1.5-1.5zM7.5 6a1.5 1.5 0 1 1-3 0a1.5 1.5 0 0 1 3 0"158 clip-rule="evenodd"159 />160 </svg>161 <div>PNG</div>162 </div>163 164 <div class="flex space-x-5">165 <!-- Buttons for Download -->166 <button167 id="copyPngBtn"168 class="w-full flex items-center justify-center space-x-2 p-3 border border-gray-300 rounded bg-gray-100"169 >170 <svg171 xmlns="http://www.w3.org/2000/svg"172 class="w-4 h-4"173 viewBox="0 0 16 16"174 >175 <path176 fill="currentColor"177 fill-rule="evenodd"178 d="M12 2.5H8A1.5 1.5 0 0 0 6.5 4v1H8a3 3 0 0 1 3 3v1.5h1A1.5 1.5 0 0 0 13.5 8V4A1.5 1.5 0 0 0 12 2.5M11 11h1a3 3 0 0 0 3-3V4a3 3 0 0 0-3-3H8a3 3 0 0 0-3 3v1H4a3 3 0 0 0-3 3v4a3 3 0 0 0 3 3h4a3 3 0 0 0 3-3zM4 6.5h4A1.5 1.5 0 0 1 9.5 8v4A1.5 1.5 0 0 1 8 13.5H4A1.5 1.5 0 0 1 2.5 12V8A1.5 1.5 0 0 1 4 6.5"179 clip-rule="evenodd"180 />181 </svg>182 <div>Copy</div>183 </button>184 <button185 id="downloadPngBtn"186 class="w-full flex items-center justify-center space-x-2 p-3 border border-gray-300 rounded bg-gray-100"187 >188 <svg189 xmlns="http://www.w3.org/2000/svg"190 class="h-4 w-4"191 viewBox="0 0 16 16"192 >193 <path194 fill="currentColor"195 fill-rule="evenodd"196 d="M6.5 6H4.029l3.499 4.276a.61.61 0 0 0 .944 0L11.971 6H9.5V2.75a.25.25 0 0 0-.25-.25h-2.5a.25.25 0 0 0-.25.25zm1.498 7.5H1.75a.75.75 0 0 0 0 1.5h12.5a.75.75 0 0 0 0-1.5zm0-1.5a2.11 2.11 0 0 1-1.631-.774L2.833 6.907A1.474 1.474 0 0 1 3.973 4.5H5V2.75C5 1.784 5.784 1 6.75 1h2.5c.966 0 1.75.784 1.75 1.75V4.5h1.026a1.474 1.474 0 0 1 1.14 2.407l-3.533 4.319c-.4.49-1 .774-1.632.774H8"197 clip-rule="evenodd"198 />199 </svg>200 <div>Download</div>201 </button>202 </div>203 </div>204 </div>205 </div>206 207 <span id="snackbar">Copied to clipboard</span>208 209 210 211 <script type="text/javascript">212 function copySvgToClipboard() {213 const svgElement = document.getElementById('svgElement');214 const svgData = new XMLSerializer().serializeToString(svgElement);215 216 // Copy the serialized SVG text to the clipboard217 navigator.clipboard.writeText(svgData).then(() => {218 var sb = document.getElementById("snackbar");219 sb.className = "show";220 setTimeout(()=>{ sb.className = sb.className.replace("show", ""); }, 3000);221 }).catch(err => {222 console.error("Could not copy SVG text to clipboard: ", err);223 });224}225 226// Function to convert SVG to PNG and copy to clipboard227function copyPngToClipboard() {228 const svgElement = document.getElementById('svgElement');229 const svgData = new XMLSerializer().serializeToString(svgElement);230 const canvas = document.createElement('canvas');231 const ctx = canvas.getContext('2d');232 const img = new Image();233 234 img.onload = function () {235 // Ensure canvas size matches the SVG236 canvas.width = svgElement.clientWidth;237 canvas.height = svgElement.clientHeight;238 239 // Clear the canvas (transparent background)240 ctx.clearRect(0, 0, canvas.width, canvas.height);241 242 // Draw the SVG image on the canvas243 ctx.drawImage(img, 0, 0);244 245 // Convert the canvas content to a PNG blob246 canvas.toBlob(blob => {247 const clipboardItem = [new ClipboardItem({ 'image/png': blob })];248 navigator.clipboard.write(clipboardItem).then(() => {249 var sb = document.getElementById("snackbar");250 sb.className = "show";251 setTimeout(()=>{ sb.className = sb.className.replace("show", ""); }, 3000);252 253 }).catch(err => {254 console.error("Could not copy PNG to clipboard: ", err);255 });256 }, 'image/png');257 };258 259 // Set the image source to the serialized SVG data260 img.src = 'data:image/svg+xml;base64,' + btoa(svgData);261}262 263// Function to download SVG264function downloadSvg() {265 const svgElement = document.getElementById('svgElement');266 const svgData = new XMLSerializer().serializeToString(svgElement);267 const blob = new Blob([svgData], { type: 'image/svg+xml' });268 const url = URL.createObjectURL(blob);269 270 const a = document.createElement('a');271 a.href = url;272 a.download = 'image.svg';273 document.body.appendChild(a);274 a.click();275 document.body.removeChild(a);276 URL.revokeObjectURL(url);277}278 279// Function to download PNG280function downloadPng() {281 const svgElement = document.getElementById('svgElement');282 const svgData = new XMLSerializer().serializeToString(svgElement);283 const canvas = document.createElement('canvas');284 const ctx = canvas.getContext('2d');285 const img = new Image();286 287 img.onload = function () {288 canvas.width = svgElement.clientWidth;289 canvas.height = svgElement.clientHeight;290 ctx.drawImage(img, 0, 0);291 canvas.toBlob(blob => {292 const url = URL.createObjectURL(blob);293 const a = document.createElement('a');294 a.href = url;295 a.download = 'image.png';296 document.body.appendChild(a);297 a.click();298 document.body.removeChild(a);299 URL.revokeObjectURL(url);300 }, 'image/png');301 };302 303 img.src = 'data:image/svg+xml;base64,' + btoa(svgData);304}305 306// Button event listeners307document.getElementById('copySvgBtn').addEventListener('click', copySvgToClipboard);308document.getElementById('copyPngBtn').addEventListener('click', copyPngToClipboard);309document.getElementById('downloadSvgBtn').addEventListener('click', downloadSvg);310document.getElementById('downloadPngBtn').addEventListener('click', downloadPng);311 </script>312 313 </body></html>"""314 )315 x = x.replace('"', """)316 x = x.replace("'", """)317 return f"""<iframe style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera; 318 display-capture; encrypted-media;" sandbox="allow-modals allow-forms 319 allow-scripts allow-same-origin allow-popups 320 allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 321 allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""322 323def predict(input_mol, style, contour_level, view_str, chains):324 # write view to file325 with open("view_matrix", "w") as f:326 f.write(json.loads(view_str))327 328 chain_str = ""329 chain_dict = json.loads(chains)330 331 outputfile = os.path.basename(input_mol).replace(".pdb",".svg")332 # sort keys in dict and add colors to chain_str333 for chain in sorted(chain_dict.keys()):334 chain_str += f" '{chain_dict[chain]}'"335 if style == "Goodsell3D":336 os.system(f"cellscape cartoon --pdb {input_mol.name} --outline residue --color_by chain --depth_shading --depth_lines --colors {chain_str} --depth flat --back_outline --view view_matrix --save {outputfile}")337 elif style == "Contour":338 os.system(f"cellscape cartoon --pdb {input_mol.name} --outline chain --color_by chain --depth_contour_interval {contour_level} --colors {chain_str} --depth contours --back_outline --view view_matrix --save {outputfile}")339 else:340 os.system(f"cellscape cartoon --pdb {input_mol.name} --outline chain --colors {chain_str} --depth flat --back_outline --view view_matrix --save {outputfile}")341 342 #read content of file343 os.system(f"inkscape {outputfile} --actions='select-all;path-simplify;export-plain-svg' --export-filename {outputfile.replace('.svg', '_opt.svg')}")344 345 346 return html_output(outputfile.replace('.svg', '_opt.svg'))347 348 349 350 351def show_contour_level(style):352 if style=="Contour":353 return gr.Slider(minimum=1,maximum=50,step=1, value=10, label="Contour level", visible=True)354 else:355 return gr.Slider(minimum=1,maximum=50,step=1, value=10, label="Contour level", visible=False)356 357with gr.Blocks() as demo:358 style = gr.Radio(value="Flat", choices=["Flat", "Contour", "Goodsell3D"], label="Style")359 contour_level = gr.Slider(minimum=1,maximum=50,step=1, value=10, label="Contour level", visible=False)360 style.change(show_contour_level, style, contour_level)361 inp = moleculeview(label="Molecule3D")362 363 view_str = gr.Textbox("viewMatrixResult", label="View Matrix", visible=False)364 chains = gr.Textbox("chainsResult", label="Chains", visible=False) 365 hidden_style = gr.Textbox(visible=False)366 timestamp = gr.Textbox(visible=False)367 368 btn = gr.Button("Vectorize")369 html = gr.HTML("")370 btn.click(None, style, [view_str, chains, hidden_style, timestamp], js="(style) => [document.getElementById('viewMatrixResult').value, document.getElementById('chains').value, style, Date.now()]") #371 372 timestamp.change(predict, [inp, style, contour_level, view_str, chains], [html])373 # on change of chains trigger, rendering374 375 376if __name__ == "__main__":377 demo.launch()378