TmMolStar/src/TmImageRenderer/index.ts

/**
 * Copyright (c) 2023 mol* contributors, licensed under MIT, See LICENSE file for more info.
 *
 * @author Adam Midlik <midlik@gmail.com>
 */
/**
 *
 * Copyright (C) 2024, Protein Bioinformatics Research Group, HUN-REN RCNS
 *
 * Licensed under CC BY-NC 4.0, see LICENSE file for more info.
 *
 * @author Gabor Tusnady <tusnady.gabor@ttk.hu>
 * @author Csongor Gerdan <gerdan.csongor@ttk.hu>
 */

import { ArgumentParser } from 'argparse';
import fs from 'fs';
import path from 'path';
import gl from 'gl';
import pngjs from 'pngjs';
import jpegjs from 'jpeg-js';

import { HeadlessPluginContext } from '../mol-plugin/headless-plugin-context';
import { DefaultPluginSpec } from '../mol-plugin/spec';
import { ExternalModules, defaultCanvas3DParams } from '../mol-plugin/util/headless-screenshot';
import { setFSModule } from '../mol-util/data-source';
import { MembraneOrientationPreset, loadWithUNITMPMembraneRepresentation } from '../tmdet-extension/behavior';
import { TmDetColorThemeProvider } from '../tmdet-extension/tmdet-color-theme';
import { TmDetLabelProvider } from '../tmdet-extension/labeling';
import { ColorNames } from '../mol-util/color/names';


const IMAGE_WIDTH = 800;
const IMAGE_HEIGHT = 800;

setFSModule(fs);

interface Args {
    pdbId: string,
    side1: string,
    outDirectory: string,
    domain: string
}

function parseArguments(): Args {
    const parser = new ArgumentParser({ description: 'Command-line application generating PDBTM images of PDB structures' });
    parser.add_argument('pdbId', { help: 'PDB identifier' });
    parser.add_argument('--side1', { default: 'Inside', help: 'Inside|Outside|Periplasm (default: Inside)' });
    parser.add_argument('outDirectory', { help: 'Directory for outputs' });
    const domain = 'http://pdbtm.unitmp.botond.hu';
    parser.add_argument('--domain', { default: domain, help: `Domain of data source site (default: ${domain})` });
    const args = parser.parse_args();
    return { ...args };
}

async function main() {
    const args = parseArguments();
    const domain = args.domain;
    const structureUrl = `${domain}/api/v1/entry/${args.pdbId}.updated.cif`;
    const regionDescriptorUrl = `${domain}/api/v1/entry/${args.pdbId}.json`;
    const fileName = `${args.pdbId}.jpg`;
    console.log('PDB ID:', args.pdbId);
    console.log('Side1:', args.side1);
    console.log('Structure URL:', structureUrl);
    console.log('Region descriptor URL:', regionDescriptorUrl);
    console.log('Outputs:', args.outDirectory);

    // Create a headless plugin
    const externalModules: ExternalModules = { gl, pngjs, 'jpeg-js': jpegjs };
    const canvasParams = defaultCanvas3DParams();
    canvasParams.camera!.mode = "perspective";
    canvasParams.renderer!.backgroundColor = ColorNames.black;
    const plugin = new HeadlessPluginContext(
        externalModules,
        DefaultPluginSpec(),
        { width: IMAGE_WIDTH, height: IMAGE_HEIGHT },
        { canvas: canvasParams }
    );
    await plugin.init();



    //
    // TMDET extension registration section
    //
    let tree = plugin.state.behaviors.build();

    plugin.builders.structure.representation.registerPreset(MembraneOrientationPreset);
    plugin.representation.structure.themes.colorThemeRegistry.add(TmDetColorThemeProvider);
    plugin.managers.lociLabels.addProvider(TmDetLabelProvider);

    await plugin.runTask(plugin.state.behaviors.updateTree(tree, { doNotUpdateCurrent: true, doNotLogTiming: true }));
    await loadWithUNITMPMembraneRepresentation(plugin, {
        structureUrl: structureUrl,
        regionDescriptorUrl: regionDescriptorUrl,
        side1: args.side1
    });
    //
    // End of TMDET extension registration section
    //



    // Export images
    fs.mkdirSync(args.outDirectory, { recursive: true });
    setTimeout(async () => {
            plugin.saveImage(path.join(args.outDirectory, fileName), { width: 100, height: 100 });
            // Cleanup
            await plugin.clear();
            plugin.dispose();
        },
        1000
    );

}

main();