Spaces:

simonduerr
/

pdb2vector

Running

File size: 14,767 Bytes


import os
import json
import uuid

import gradio as gr
from gradio_moleculeview import moleculeview


import cellscape


def html_output(input_file):
    with open(input_file, "r") as f:
        svg = f.read().replace("<svg", "<svg id='svgElement'")
        
    x = (
        """<!DOCTYPE html>
<html lang="en">
  <head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>Document</title>
    <script src="https://cdn.tailwindcss.com"></script>

    <style>
    #snackbar {
  visibility: hidden;
  color: #fff;
  background-color: #333;
  min-width: 250px;
  margin-left: -125px;
  border-radius: 2px;
  padding: 16px;
  text-align: center;
  left: 50%;
  bottom: 30px;
  z-index: 1;
  position: fixed;
}

/* This will be activated when the snackbar's class is 'show' which will be added through JS */
#snackbar.show {
  visibility: visible;
  -webkit-animation: fadein 0.5s, fadeout 0.5s 2.5s;
  animation: fadein 0.5s, fadeout 0.5s 2.5s;
}

/* Animations for fading in and out */
@-webkit-keyframes fadein {
  from {bottom: 0; opacity: 0;}
  to {bottom: 30px; opacity: 1;}
}

@keyframes fadein {
  from {bottom: 0; opacity: 0;}
  to {bottom: 30px; opacity: 1;}
}

@-webkit-keyframes fadeout {
  from {bottom: 30px; opacity: 1;}
  to {bottom: 0; opacity: 0;}
}

@keyframes fadeout {
  from {bottom: 30px; opacity: 1;}
  to {bottom: 0; opacity: 0;}
}
    </style>
  </head>
  <body>

    <div class="flex container mx-auto max-w-3xl">
      <div class="w-1/2">
        
     

 
    """
        + svg
        + """
         </div>

      <div class="w-1/2 space-y-4 justify-center items-center flex flex-col">
        <div class="w-3/4">
          <div
            class="w-full flex space-x-2 items-center text-xl justify-center font-light mb-3"
          >
            <svg
              xmlns="http://www.w3.org/2000/svg"
              class="w-4 h-4"
              viewBox="0 0 16 16"
            >
              <path
                fill="currentColor"
                fill-rule="evenodd"
                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"
                clip-rule="evenodd"
              />
            </svg>
            <div>SVG</div>
          </div>
          <div class="w-full flex space-x-5">
            <button
              id="copySvgBtn"
              class="flex items-center justify-center space-x-2 w-full p-3 border border-gray-300 rounded bg-gray-100"
            >
              <svg
                xmlns="http://www.w3.org/2000/svg"
                class="w-4 h-4"
                viewBox="0 0 16 16"
              >
                <path
                  fill="currentColor"
                  fill-rule="evenodd"
                  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"
                  clip-rule="evenodd"
                />
              </svg>
              <div>Copy*</div>
            </button>
            <button
              id="downloadSvgBtn"
              class="w-full flex items-center justify-center space-x-2 p-3 border border-gray-300 rounded bg-gray-100"
            >
              <svg
                xmlns="http://www.w3.org/2000/svg"
                class="h-4 w-4"
                viewBox="0 0 16 16"
              >
                <path
                  fill="currentColor"
                  fill-rule="evenodd"
                  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"
                  clip-rule="evenodd"
                />
              </svg>
              <div>Download</div>
            </button>
          </div>

          <div class="text-xs mt-3 text-center w-full">
            *Only works in in Adobe Illustrator.
          </div>

          <div
            class="flex space-x-2 items-center text-xl justify-center font-light mb-3 mt-10"
          >
            <svg
              xmlns="http://www.w3.org/2000/svg"
              class="w-4 h-4"
              viewBox="0 0 16 16"
            >
              <path
                fill="currentColor"
                fill-rule="evenodd"
                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"
                clip-rule="evenodd"
              />
            </svg>
            <div>PNG</div>
          </div>

          <div class="flex space-x-5">
            <!-- Buttons for Download -->
            <button
              id="copyPngBtn"
              class="w-full flex items-center justify-center space-x-2 p-3 border border-gray-300 rounded bg-gray-100"
            >
              <svg
                xmlns="http://www.w3.org/2000/svg"
                class="w-4 h-4"
                viewBox="0 0 16 16"
              >
                <path
                  fill="currentColor"
                  fill-rule="evenodd"
                  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"
                  clip-rule="evenodd"
                />
              </svg>
              <div>Copy</div>
            </button>
            <button
              id="downloadPngBtn"
              class="w-full flex items-center justify-center space-x-2 p-3 border border-gray-300 rounded bg-gray-100"
            >
              <svg
                xmlns="http://www.w3.org/2000/svg"
                class="h-4 w-4"
                viewBox="0 0 16 16"
              >
                <path
                  fill="currentColor"
                  fill-rule="evenodd"
                  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"
                  clip-rule="evenodd"
                />
              </svg>
              <div>Download</div>
            </button>
          </div>
        </div>
      </div>
    </div>

  <span id="snackbar">Copied to clipboard</span>


      
    <script type="text/javascript">
    function copySvgToClipboard() {
    const svgElement = document.getElementById('svgElement');
    const svgData = new XMLSerializer().serializeToString(svgElement);

    // Copy the serialized SVG text to the clipboard
    navigator.clipboard.writeText(svgData).then(() => {
     var sb = document.getElementById("snackbar");
              sb.className = "show";
              setTimeout(()=>{ sb.className = sb.className.replace("show", ""); }, 3000);
    }).catch(err => {
        console.error("Could not copy SVG text to clipboard: ", err);
    });
}

// Function to convert SVG to PNG and copy to clipboard
function copyPngToClipboard() {
    const svgElement = document.getElementById('svgElement');
    const svgData = new XMLSerializer().serializeToString(svgElement);
    const canvas = document.createElement('canvas');
    const ctx = canvas.getContext('2d');
    const img = new Image();

    img.onload = function () {
        // Ensure canvas size matches the SVG
        canvas.width = svgElement.clientWidth;
        canvas.height = svgElement.clientHeight;

        // Clear the canvas (transparent background)
        ctx.clearRect(0, 0, canvas.width, canvas.height);

        // Draw the SVG image on the canvas
        ctx.drawImage(img, 0, 0);

        // Convert the canvas content to a PNG blob
        canvas.toBlob(blob => {
            const clipboardItem = [new ClipboardItem({ 'image/png': blob })];
            navigator.clipboard.write(clipboardItem).then(() => {
             var sb = document.getElementById("snackbar");
              sb.className = "show";
              setTimeout(()=>{ sb.className = sb.className.replace("show", ""); }, 3000);
  
            }).catch(err => {
                console.error("Could not copy PNG to clipboard: ", err);
            });
        }, 'image/png');
    };

    // Set the image source to the serialized SVG data
    img.src = 'data:image/svg+xml;base64,' + btoa(svgData);
}

// Function to download SVG
function downloadSvg() {
    const svgElement = document.getElementById('svgElement');
    const svgData = new XMLSerializer().serializeToString(svgElement);
    const blob = new Blob([svgData], { type: 'image/svg+xml' });
    const url = URL.createObjectURL(blob);

    const a = document.createElement('a');
    a.href = url;
    a.download = 'image.svg';
    document.body.appendChild(a);
    a.click();
    document.body.removeChild(a);
    URL.revokeObjectURL(url);
}

// Function to download PNG
function downloadPng() {
    const svgElement = document.getElementById('svgElement');
    const svgData = new XMLSerializer().serializeToString(svgElement);
    const canvas = document.createElement('canvas');
    const ctx = canvas.getContext('2d');
    const img = new Image();

    img.onload = function () {
        canvas.width = svgElement.clientWidth;
        canvas.height = svgElement.clientHeight;
        ctx.drawImage(img, 0, 0);
        canvas.toBlob(blob => {
            const url = URL.createObjectURL(blob);
            const a = document.createElement('a');
            a.href = url;
            a.download = 'image.png';
            document.body.appendChild(a);
            a.click();
            document.body.removeChild(a);
            URL.revokeObjectURL(url);
        }, 'image/png');
    };

    img.src = 'data:image/svg+xml;base64,' + btoa(svgData);
}

// Button event listeners
document.getElementById('copySvgBtn').addEventListener('click', copySvgToClipboard);
document.getElementById('copyPngBtn').addEventListener('click', copyPngToClipboard);
document.getElementById('downloadSvgBtn').addEventListener('click', downloadSvg);
document.getElementById('downloadPngBtn').addEventListener('click', downloadPng);
    </script>
           
        </body></html>"""
    )
    x = x.replace('"', "&quot;")
    x = x.replace("'", "&quot;")
    return f"""<iframe style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera; 
    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
    allow-scripts allow-same-origin allow-popups 
    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
    allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""

def predict(input_mol, style, contour_level, view_str, chains):
    # write view to file
    with open("view_matrix", "w") as f:
        f.write(json.loads(view_str))

    chain_str = ""
    chain_dict = json.loads(chains)

    outputfile = os.path.basename(input_mol).replace(".pdb",".svg")
    # sort keys in dict and add colors to chain_str
    for chain in sorted(chain_dict.keys()):
        chain_str += f" '{chain_dict[chain]}'"
    if style == "Goodsell3D":
        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}")
    elif style == "Contour":
        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}")
    else:
        os.system(f"cellscape cartoon --pdb {input_mol.name} --outline chain --colors {chain_str}  --depth flat --back_outline --view view_matrix --save {outputfile}")
    
    #read content of file
    os.system(f"inkscape {outputfile} --actions='select-all;path-simplify;export-plain-svg' --export-filename {outputfile.replace('.svg', '_opt.svg')}")

    
    return html_output(outputfile.replace('.svg', '_opt.svg'))




def show_contour_level(style):
    if style=="Contour":
        return gr.Slider(minimum=1,maximum=50,step=1, value=10, label="Contour level", visible=True)
    else:
        return gr.Slider(minimum=1,maximum=50,step=1, value=10, label="Contour level", visible=False)

with gr.Blocks() as demo:
    style = gr.Radio(value="Flat", choices=["Flat", "Contour", "Goodsell3D"], label="Style")
    contour_level = gr.Slider(minimum=1,maximum=50,step=1, value=10, label="Contour level", visible=False)
    style.change(show_contour_level, style, contour_level)
    inp = moleculeview(label="Molecule3D")

    view_str = gr.Textbox("viewMatrixResult", label="View Matrix", visible=False)
    chains = gr.Textbox("chainsResult", label="Chains", visible=False) 
    hidden_style = gr.Textbox(visible=False)
    timestamp = gr.Textbox(visible=False)
    
    btn = gr.Button("Vectorize")
    html = gr.HTML("")
    btn.click(None, style, [view_str, chains, hidden_style, timestamp], js="(style) => [document.getElementById('viewMatrixResult').value, document.getElementById('chains').value, style, Date.now()]")  #
    
    timestamp.change(predict, [inp, style, contour_level, view_str, chains], [html])
    # on change of chains trigger, rendering
    

if __name__ == "__main__":
    demo.launch()