dtm_to_heightmap2.py

#!/usr/bin/env python

import os
import sys
import argparse
import logging
import numpy
from planetaryimage.pds3image import PDS3Image
from struct import pack, unpack

class DtmToHeightmap :
    DTM_FileName = ''
    Photo_FileName = ''

    DTM_Image = None
    Photo_Image = None

    def __init__(self, DTM_FileName, Photo_FileName, Latitude, Longitude, OutputDir):

        self.DTM_FileName = DTM_FileName
        self.Photo_FileName = Photo_FileName
        self.CenterLatitude = Latitude
        self.CenterLongitude = Longitude
        self.OutputDir = OutputDir

        logging.info("Reading {}".format(self.DTM_FileName))
        self.DTM_Image = PDS3Image.open(self.DTM_FileName)
        logging.info("Reading {}".format(self.Photo_FileName))
        self.Photo_Image = PDS3Image.open(self.Photo_FileName)

        self.HeightMapSize = 257 # 257 pixel x 257 pixel (must be 2^n+1)
        (self.DTM_CroppedImage, self.DTM_OutputName)  = self.cropped_image(self.DTM_Image, self.HeightMapSize)
        self.write_dtm_image()

        #self.TextureMapSize = 512 # 512 pixel x 512 pixel (must be 2^n)
        self.TextureMapSize = 1024 # 1024 pixel x 1024 pixel (must be 2^n)
        (self.Photo_CroppedImage, self.Photo_OutputName)  = self.cropped_image(self.Photo_Image, self.TextureMapSize)
        self.write_photo_image()
        self.write_world_files()

    def cropped_image(self, Image, MapSize):

        logging.debug(Image.label)
        l = Image.label['IMAGE_MAP_PROJECTION']
        MapResolution = l['MAP_RESOLUTION'].value
        MapScale = l['MAP_SCALE'].value
        MaximumLatitude = l['MAXIMUM_LATITUDE'].value
        MinimumLatitude = l['MINIMUM_LATITUDE'].value
        EasternmostLongitude = l['EASTERNMOST_LONGITUDE'].value
        WesternmostLongitude = l['WESTERNMOST_LONGITUDE'].value
        l = Image.label['IMAGE']
        LinePixel = l['LINES']
        SamplePixel = l['LINE_SAMPLES']

        logging.info("..... Map Resolution {} <PIX/DEG>".format(MapResolution))
        logging.info("..... Map Sale {} <METERS>".format(MapScale))
        logging.info("..... Maximum Latitude {} <DEG>".format(MaximumLatitude))
        logging.info("..... Minimum Latitude {} <DEG>".format(MinimumLatitude))
        logging.info("..... Easternmost Longitude {} <DEG>".format(EasternmostLongitude))
        logging.info("..... Westernmost Longitude {} <DEG>".format(WesternmostLongitude))
        logging.info("..... {} pixel / line".format(LinePixel))
        logging.info("..... {} line / sample".format(SamplePixel))

        StartLatitude = self.CenterLatitude + MapSize/2/MapResolution
        StartLongitude =  self.CenterLongitude - MapSize/2/MapResolution
        EndLatitude = self.CenterLatitude - MapSize/2/MapResolution
        EndLongitude =  self.CenterLongitude + MapSize/2/MapResolution
        print("HightMap Cropping Area {} - {}".format(StartLatitude, EndLatitude))
        print("                       {} - {}".format(StartLongitude, EndLongitude))
        if StartLatitude > MaximumLatitude or StartLatitude < MinimumLatitude or StartLongitude < WesternmostLongitude or EndLongitude > EasternmostLongitude:
            print("ERROR: Cropping  outside of the map")
            exit(1)

        OffsetLatitude = int(round((MaximumLatitude - StartLatitude)*MapResolution))
        OffsetLongitude = int(round((StartLongitude - WesternmostLongitude)*MapResolution))

        print("HightMap Cropping Pixel Latitude  {} - {} - {}".format(OffsetLatitude, OffsetLatitude+MapSize/2, OffsetLatitude+MapSize))
        print("                        Longitude {} - {} - {}".format(OffsetLongitude, OffsetLongitude+MapSize/2, OffsetLongitude+MapSize) )
        
        return (Image.image[
            OffsetLatitude:OffsetLatitude+MapSize, 
            OffsetLongitude:OffsetLongitude+MapSize
        ],
                "{}_{}x{}+{}+{}.png".format(Image.label['PRODUCT_ID'], MapSize, MapSize, OffsetLatitude, OffsetLongitude))

    def write_dtm_image(self):
        
        SkipData = [pack('BBBB', 0xFB, 0xFF, 0x7F, 0xFF), pack('BBBB', 0xFC, 0xFF, 0x7F, 0xFF), pack('BBBB', 0xFD, 0xFF, 0x7F, 0xFF), pack('BBBB', 0xFF, 0xFF, 0x7F, 0xFF), pack('BBBB', 0xFE, 0xFF, 0x7F, 0xFF)]
        self.MinimumAltitude = min(filter(lambda n: not pack('f', n) in SkipData, self.DTM_CroppedImage.flatten()))
        self.MaximumAltitude = max(filter(lambda n: not pack('f', n) in SkipData, self.DTM_CroppedImage.flatten()))
        NumberOfSkipedData = len(filter(lambda n: pack('f', n) in SkipData, self.DTM_CroppedImage.flatten()))

        logging.info("..... Minimum Altitude MinimumAltitude {}".format(self.MinimumAltitude))
        logging.info("..... Maximul Altitude MaximumAltitude {}".format(self.MaximumAltitude))

        if float(NumberOfSkipedData)/(self.HeightMapSize*self.HeightMapSize) > 0.01:
            logging.info("--- Assuming there is a pit")
            self.MinimumAltitude -= 100 # pit depth is about 100[m]? http://lroc.sese.asu.edu/posts/230
        HeightResolution = (self.MaximumAltitude - self.MinimumAltitude)/255
        self.HeightCenter = -(self.DTM_CroppedImage[self.HeightMapSize/2][self.HeightMapSize/2]-self.MinimumAltitude)

        logging.info("..... Altitude resolution {}".format(HeightResolution))
        logging.info("..... Skipped Pixel {0} ({1} %)".format(NumberOfSkipedData, 100*NumberOfSkipedData/(self.HeightMapSize*self.HeightMapSize)))

        for j in range(self.HeightMapSize):
            for i in range(self.HeightMapSize):
                if self.MinimumAltitude < self.DTM_CroppedImage[i][j] and self.DTM_CroppedImage[i][j] < self.MaximumAltitude:
                    self.DTM_CroppedImage[i][j] = int((self.DTM_CroppedImage[i][j]-self.MinimumAltitude)/HeightResolution)
                else:
                    self.DTM_CroppedImage[i][j] = 0 # pit

        # write
        if not os.path.exists(os.path.join(self.OutputDir,"media/materials/textures/")):
            os.makedirs(os.path.join(self.OutputDir,"media/materials/textures/"))

        logging.info("--- Output {}".format(self.DTM_OutputName))
        # curl -LO https://raw.github.com/drj11/pypng/master/code/png.py
        import png
        f = open(os.path.join(self.OutputDir, "media/materials/textures/", self.DTM_OutputName), 'wb')
        w = png.Writer(self.HeightMapSize, self.HeightMapSize, greyscale=True)
        w.write(f, self.DTM_CroppedImage)
        f.close()

    def write_photo_image(self):
        
        # write
        if not os.path.exists(os.path.join(self.OutputDir,"media/materials/textures/")):
            os.makedirs(os.path.join(self.OutputDir,"media/materials/textures/"))

        for j in range(self.TextureMapSize):
            for i in range(self.TextureMapSize):
                self.Photo_CroppedImage[j][i] = self.Photo_CroppedImage[j][i]/(self.Photo_Image.label['IMAGE']['SAMPLE_BITS']/8)
                if self.Photo_CroppedImage[j][i] < 0:
                    self.Photo_CroppedImage[j][i] = 0
                if self.Photo_CroppedImage[j][i] >= 255:
                    self.Photo_CroppedImage[j][i] = 255

        logging.info("--- Output {}".format(self.Photo_OutputName))
        # curl -LO https://raw.github.com/drj11/pypng/master/code/png.py
        import png
        f = open(os.path.join(self.OutputDir, "media/materials/textures/", self.Photo_OutputName), 'wb')
        w = png.Writer(self.TextureMapSize, self.TextureMapSize, greyscale=True)
        w.write(f, self.Photo_CroppedImage)
        f.close()

    def write_world_files(self):

        # model.config
        fname = os.path.join(self.OutputDir,"model.config")
        logging.info("--- Output {}".format(fname))
        f = open(fname, 'w')
        f.write('''
<?xml version="1.0"?>
<model>
  <name>{name}</name>
  <version>1.0</version>
  <sdf version="1.4">model-1_4.sdf</sdf>
  <sdf version="1.5">model.sdf</sdf>

  <author>
    <name>DTM to SDF File Converter</name>
    <email>k-okada@jsk.t.u-tokyo.ac.jp</email>
  </author>

  <description>
{command}

{description}
  </description>
</model>
        '''.format(name=self.OutputDir,command=' '.join(sys.argv), description=str(self.DTM_Image.label)))

        # model.sdf
        MapScale = self.DTM_Image.label['IMAGE_MAP_PROJECTION']['MAP_SCALE'].value

        fname = os.path.join(self.OutputDir,"model.sdf")
        logging.info("--- Output {}".format(fname))
        f = open(fname, 'w')
        model_sdf = '''
<?xml version="1.0" ?>
<sdf version="{{version}}">
  <model name="{name}">
    <static>true</static>
    <link name="link">
      <collision name="collision">
        <surface>
          <friction>
          <ode>
            <mu>0.2</mu>
          </ode>
          </friction>
          <contact>
          <ode>
            <soft_cfm>1</soft_cfm>
            <kp>100000</kp>
            <kd>1</kd>
            <max_vel>0.000001</max_vel>
            <min_depth>0.02</min_depth>
          </ode>
          </contact>
        </surface>
        <geometry>
          <heightmap>
            <uri>model://{name}/media/materials/textures/{heightmap}</uri>
            <size>{size} {size} {height}</size>
            <pos>0 0 {height_center}</pos>
          </heightmap>
        </geometry>
      </collision>

      <visual name="visual">
        <geometry>
          <heightmap>
            <texture>
              <diffuse>file://media/materials/textures/{texture}</diffuse>
              <normal>file://media/materials/textures/flat_normal.png</normal>
              <size>{size}</size>
            </texture>
            <uri>model://{name}/media/materials/textures/{heightmap}</uri>
            <size>{size} {size} {height}</size>
            <pos>0 0 {height_center}</pos>
          </heightmap>
        </geometry>
      </visual>
    </link>
  </model>
</sdf>
        '''.format(name=self.OutputDir, heightmap=self.DTM_OutputName, texture=self.Photo_OutputName, size=int(self.HeightMapSize*MapScale), height=int(self.MaximumAltitude - self.MinimumAltitude), height_center=int(self.HeightCenter))
        f.write(model_sdf.format(version=1.5))

        fname = os.path.join(self.OutputDir,"model-1_4.sdf")
        logging.info("--- Output {}".format(fname))
        f = open(fname, 'w')
        f.write(model_sdf.format(version=1.4))

        logging.info("--- Done writing")

# script to parse dtm file into gazebo world file
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='This script convert DTM file to gazebo world file, for example http://lroc.sese.asu.edu/data/LRO-L-LROC-5-RDR-V1.0/LROLRC_2001/DATA/SDP/NAC_DTM/TRANQPIT1/NAC_DTM_TRANQPIT1_E080N0330.IMG')
    parser.add_argument('DTM_FileName', action='store', type=str, help='Digital Terrain Model (PDS IMG) file name')
    parser.add_argument('Photo_FileName', action='store', type=str, help='Orthophoto (TIFF 0.50m/px) file name')
    parser.add_argument('--output-dir', action='store', type=str, help='Output direcotry name')
    parser.add_argument('--latitude', action='store', type=float, help='Latitude of the center of the map')
    parser.add_argument('--longitude', action='store', type=float, help='Longitude of the center of the map')
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO)

    DtmToHeightmap(DTM_FileName=args.DTM_FileName, Photo_FileName=args.Photo_FileName, Latitude=args.latitude, Longitude=args.longitude, OutputDir=args.output_dir)