orthographic/photogen.py

import os
from PIL import Image, ImageFilter, ImageEnhance, ImageFile

from defines import *
from layergen import *
from log import *
from functions import *
from xp_normalmap import *

ImageFile.LOAD_TRUNCATED_IMAGES = True

# -------------------------------------------------------------------
# ORTHOGRAPHIC
# Your personal aerial satellite. Always on. At any altitude.*
# Developed by MarStrMind
# License: Open Software License 3.0
# Up to date version always on marstr.online
# -------------------------------------------------------------------
# photogen.py
# The class that generates the photo tiles from previous layers,
# in their correct order.
# -------------------------------------------------------------------


class mstr_photogen:

    # Initializer doesn't need much
    def __init__ (self, lat, lng, ty, tx, maxlat, maxlng):
        self._lat = lat
        self._lng = lng
        self._ty = ty
        self._tx = tx
        self._maxlatlng = [ maxlat, maxlng ]
        # Define layer size depending on what is wanted
        self._imgsize = mstr_photores
        # Empty image where everything goes into
        self._tile = Image.new("RGBA", (self._imgsize, self._imgsize))
        self._latlngfld = self.latlng_folder([lat,lng])
        mstr_msg("photogen", "Photogen initialized")


    # Defines the order of layer names that were processed
    # Called by orthographic prior to starting the photo gen process
    def setLayerNames(self, names):
        self._lyrnames = names

    
    # This puts it all together. Bonus: AND saves it.
    def genphoto(self, layers, waterlayers, cpl):
        # Template for the file name which is always the same
        #root_filename = mstr_datafolder + "/_cache/" + str(self._lat) + "-" + str(self._ty) + "_" + str(self._lng) + "-" + str(self._tx) + "_"

        # Correct layers
        #mstr_msg("photogen", "Correcting layer order issues")
        #layers = self.correctLayerIssues(layers)

        # First, we walk through all layers and blend them on top of each other, in order
        mstr_msg("photogen", "Merging layers")

        lyr=0
        for l in layers:
            if self._lyrnames[lyr] != "building":
                self._tile.alpha_composite(l)
            lyr=lyr+1


        # When we have run through this loop, we will end up with a sandwiched
        # image of all the other images, in their correct order.
        # However, since I have discovered that some areas in OSM simply do not
        # have any tag or information, it is possible that the final image will
        # have empty, alpha-transparent patches.
        # For this reason we need to check against these and fix that.
        
        # First, we will check if there is something to fix:
        emptyspace = self.checkForEmptySpace()
        mstr_msg("photogen", "Checked for empty patches")

        # If this check comes back as true, we need to perform
        # aforementioned fix:
        if emptyspace == True:

            mstr_msg("photogen", "Patching empty space")

            cmpl = Image.new("RGBA", (self._imgsize, self._imgsize))

            edn = self.find_earthnavdata_number()
            edns = self.latlng_folder(edn)

            cplstr = ""
            for c in range(0, len(cpl)):
                cplstr = cplstr + str(cpl[c])
                if c < len(cpl)-1:
                    cplstr = cplstr + "_"

            # Find the right color catalogue
            cpl_catalog = 0
            for c in range(len(mstr_completion_colors)):
                if mstr_completion_colors[c][0] == edn:
                    cpl_catalog = c

            # Put in some pixels
            cmpl_bg = Image.new("RGBA", (self._tile.width, self._tile.height))
            cmpl_pix = cmpl_bg.load()
            for y in range(0, self._tile.height):
                for x in range(0, self._tile.width):
                    idx = randrange(0, len(mstr_completion_colors[cpl_catalog][1]))
                    clr = mstr_completion_colors[cpl_catalog][1][idx]
                    cmpl_pix[x,y] = clr
            cmpl_bg = ImageEnhance.Contrast(cmpl_bg).enhance(0.8)
            cmpl_bg = cmpl_bg.filter(ImageFilter.GaussianBlur(radius=1))
            cmpl_bg.alpha_composite(self._tile)
            self._tile = cmpl_bg

            # Patches to add from other sources. If they don't exist, we also need to make them
            masks = glob.glob(mstr_datafolder + "textures/tile/completion/*.png")
            amt = randrange(5, 16)
            patchtags = [
                ["landuse", "meadow"],
                ["landuse", "grass"],
                ["natural", "heath"],
                ["natural", "scrub"]
            ]

            for i in range(1, amt + 1):
                pick = randrange(0, len(masks))
                patchmask = Image.open(masks[pick])
                patchmask = patchmask.rotate(randrange(0, 360), expand=True)
                patchpix = patchmask.load()

                # Pick from possible tags and values for the patches
                numbers = list(range(1, 16))
                src = random.sample(numbers, 5)

                patchpick = randrange(0, len(patchtags))

                rg = mstr_resourcegen(patchtags[patchpick][0], patchtags[patchpick][1], src)
                rg.setLayerContrast(randrange(1, 4))
                ptch = rg.gensource()

                rg_img = ptch[0]
                rg_pix = rg_img.load()

                # The patch to be used in the layer
                layerpatch = Image.new("RGBA", (patchmask.width, patchmask.height))
                lp_pix = layerpatch.load()
                for y in range(0, patchmask.height):
                    for x in range(0, patchmask.width):
                        ptc_msk = patchpix[x, y]
                        if ptc_msk[3] > 0:
                            oc = rg_pix[x, y]
                            nc = (oc[0], oc[1], oc[2], ptc_msk[3])
                            lp_pix[x, y] = nc

                #layerpatch = layerpatch.rotate(randrange(0, 360), expand=True)

                lx = randrange(0, self._imgsize - layerpatch.width)
                ly = randrange(0, self._imgsize - layerpatch.height)
                cmpl.alpha_composite(layerpatch, (lx, ly))

            # Merge the images
            cmpl.alpha_composite(self._tile)

            # Make this the real one
            self._tile = cmpl


        # There may be some tiles that have a larger sea or even an ocean in them - these need to be
        # removed from the final tile
        ocean_pix = self._tile.load()
        for y in range(self._tile.width):
            for x in range(self._tile.height):
                p = ocean_pix[x,y]
                if p[0] == 255 and p[1] == 0 and p[2] == 255:
                    t = (0,0,0,0)
                    ocean_pix[x,y] = t

        # Alpha correction on final image
        corrpix = self._tile.load()
        for y in range(0, self._tile.height):
            for x in range(0, self._tile.width):
                c = corrpix[x,y]
                if c[3] > 0:
                    nc = (c[0], c[1], c[2], 255)
                    corrpix[x,y] = nc
                if c[3] == 0:
                    corrpix[x,y] = (0,0,0,0)

        # One more thing...
        mstr_msg("photogen", "Adding features to layers")
        self.addTreesToFeatures(layers)

        # Throw missing buildings on top
        lyr = 0
        for l in layers:
            if self._lyrnames[lyr] == "building":
                self._tile.alpha_composite(l)
            lyr = lyr + 1
        
        # We are now in posession of the final image.
        
        # Contrast
        #self._tile = ImageEnhance.Contrast(self._tile).enhance(0.1)

        # This we can save accordingly.
        self._tile.save(mstr_datafolder + "z_orthographic/orthos/" + self._latlngfld + "/" + str(self._ty) + "_" + str(self._tx) + ".png")

        # Now we convert this into a DDS
        _tmpfn = mstr_datafolder + "z_orthographic/orthos/" + self._latlngfld + "/" + str(self._ty) + "_" + str(self._tx)
        os.system(mstr_xp_ddstool + " --png2dxt1 " + _tmpfn + ".png " + _tmpfn + ".dds" )
        
        os.remove(mstr_datafolder + "z_orthographic/orthos/" + self._latlngfld + "/" + str(self._ty) + "_" + str(self._tx) + ".png")


        # Now generate the normal map for this ortho.
        # But only if this is enabled.
        if mstr_xp_genscenery and mstr_xp_scn_normalmaps:
            # Generate the normal normal map first (hah)
            nrm = mstr_xp_normalmap()
            nrmimg = nrm.generate_normal_map_for_layer(self._tile, False)

            # Now we need to walk through the water layers and generate a combined normal map
            wtrlyr = Image.new("RGBA", (self._imgsize, self._imgsize))
            for w in waterlayers:
                wtrlyr.alpha_composite(w)
            wtrlyr = wtrlyr.resize((int(mstr_photores/4), int(mstr_photores/4)), Image.Resampling.BILINEAR)
            wtrimg = nrm.generate_normal_map_for_layer(wtrlyr, True)

            # Blend
            nrmimg.alpha_composite(wtrimg)

            # Save
            nrmfln = mstr_datafolder + "z_orthographic/normals/" + self._latlngfld + "/" + str(self._ty) + "_" + str(
                self._tx) + ".png"
            nrmimg.save(nrmfln)


    # Generates some random tree.
    # We will now move away from using pre-made trees...
    # they didn't look so great
    def generate_tree(self):
        sx = randrange(18, 31)
        sy = randrange(18, 31)

        treepoly = Image.new("RGBA", (sx, sy))
        draw = ImageDraw.Draw(treepoly)

        draw.ellipse((4, 4, sx - 4, sy - 4), fill="black")

        tree = Image.new("RGBA", (sx, sy))
        treepx = tree.load()
        maskpx = treepoly.load()

        # How many tree points do we want?
        treepts = 75
        # How many of those have been drawn?
        ptsdrawn = 0

        bc = [
            (36, 50, 52),
            (30, 41, 39),
            (32, 45, 37),
            (32, 39, 49),
            (33, 34, 40),
            (44, 50, 53),
            (40, 46, 48),
            (14, 31, 38),
            (17, 41, 33),
            (39, 56, 35),
            (51, 51, 42),
            (12, 27, 31),
            (45, 59, 48),
            (37, 54, 29),
            (59, 50, 34),
            (59, 59, 35),
            (59, 51, 35),
            (70, 72, 45),
            (48, 59, 44),
            (29, 47, 23),
            (47, 61, 43),
            (29, 68, 15),
            (53, 77, 63),
            (20, 68, 40)
        ]

        bcp = randrange(0, len(bc))

        treedraw = ImageDraw.Draw(tree)
        while ptsdrawn < treepts + 1:
            rx = randrange(0, sx)
            ry = randrange(0, sy)
            mp = maskpx[rx, ry]
            if mp[3] > 0:
                d = randrange(0, 51)
                r = bc[bcp][0]
                g = bc[bcp][1] + d
                b = bc[bcp][2] + (d // 2)
                a = randrange(170, 256)
                c = (r, g, b, a)
                ex = randrange(2, 5)
                ey = randrange(2, 5)
                treedraw.ellipse((rx - (ex // 2), ry - (ey // 2), rx + (ey // 2), ry + (ey // 2)), fill=c)
                ptsdrawn = ptsdrawn + 1

        for y in range(0, tree.height):
            for x in range(0, tree.width):
                tp = treepx[x, y]
                diff = randrange(0, 31)
                nc = (tp[0] - diff, tp[1] - diff, tp[2] - diff, tp[3])
                treepx[x, y] = nc

        tree = tree.filter(ImageFilter.GaussianBlur(radius=0.5))

        for y in range(0, tree.height):
            for x in range(0, tree.width):
                tp = treepx[x, y]
                diff = randrange(0, 51)
                nc = (tp[0] - diff, tp[1] - diff, tp[2] - diff, tp[3])
                treepx[x, y] = nc

        return tree


    # This used to be in layergen and solves the problem of roads being rendered above trees.
    # It is the only solution that worked, after some research.
    def addTreesToFeatures(self, layers):

        # Walk through list of layers to decide where to add the trees
        curlyr = 0
        for lyr in self._lyrnames:
            # Add trees only in some features
            if (lyr[0] == "landuse" and lyr[1] == "cemetery") or (
                    lyr[0] == "landuse" and lyr[1] == "residential") or (
                    lyr[0] == "leisure" and lyr[1] == "park"):
                trees = Image.new("RGBA", (self._imgsize, self._imgsize))
                amt = 4000
                lyrmask = layers[curlyr].load() # We can use the layer image as alpha mask
                for i in range(1, amt + 1):
                    lx = randrange(0, self._imgsize)
                    ly = randrange(0, self._imgsize)
                    lp = lyrmask[lx,ly]
                    if lp[3] > 0:
                        tree = self.generate_tree()
                        trees.alpha_composite(tree, (lx, ly))

                tree_shadow = Image.new("RGBA", (self._imgsize, self._imgsize))
                tree_pix = trees.load()
                shadow_pix = tree_shadow.load()
                for y in range(self._imgsize):
                    for x in range(self._imgsize):
                        tp = tree_pix[x, y]
                        if tp[3] > 0:
                            rndshd = randrange(5, 210)
                            sc = (0, 0, 0, rndshd)
                            if x + 8 < self._imgsize and y + 5 < self._imgsize:
                                shadow_pix[x + 8, y + 5] = sc
                tree_shadow = tree_shadow.filter(ImageFilter.GaussianBlur(radius=2))
                tree_shadow.alpha_composite(trees)
                self._tile.alpha_composite(tree_shadow)
            curlyr = curlyr + 1

        # Reset
        curlyr = 0
        bldg = []
        tilepix = self._tile.load()
        for lyr in self._lyrnames:
            if lyr[0] == "building":
                bldg.append(curlyr)
            curlyr = curlyr + 1

        for b in range(0, len(bldg)):
            bldg_lyr = layers[bldg[b]].load()
            shdw = Image.open(mstr_datafolder + "_cache/" + str(self._lat) + "-" + str(self._ty) + "_" + str(self._lng) + "-" + str(self._tx) + "_" + self._lyrnames[bldg[b]][0] + "-" + self._lyrnames[bldg[b]][1] + "_layer_shadow.png")
            shdw_pix = shdw.load()
            for y in range(0, self._imgsize):
                for x in range(0, self._imgsize):
                    bpix = bldg_lyr[x,y]
                    spix = shdw_pix[x,y]
                    if bpix[3] > 0 and spix[0] == 255:
                        tilepix[x,y] = ( bpix[0], bpix[1], bpix[2], bpix[3] )



    # Correct some layer issues
    def correctLayerIssues(self, layers):

        # First the residential/forest dilemma
        residential = 0
        forest = 0

        curlyr = 0
        for lyr in self._lyrnames:
            if lyr[0] == "landuse" and lyr[1] == "residential":
                residential=curlyr
            if lyr[0] == "landuse" and lyr[1] == "forest":
                forest = curlyr
            curlyr = curlyr+1

        layers[forest].alpha_composite(layers[residential])

        return layers




    # This checks the final image for empty patches. Should one be
    # found, we will generate something to fill the gap. If this is
    # the case, we will also note this in the database for the tile,
    # under the special tag and value "tile", "completion". The same
    # conditions apply for edge testing and so on.
    def checkForEmptySpace(self):
        empty = False

        # Load photo
        layer_pix = self._tile.load() 
        
        # Scan!
        for y in range(self._tile.width-1):
            for x in range(self._tile.height-1):
                p = layer_pix[x,y]
                if p[3] < 255: # <- Check for empty or non-complete alpha
                    empty = True
                    break

        # Tell about findings
        return empty
    

    # This returns a mask of the empty space to cover, should there be any
    def buildCompletionMask(self):
        mask = Image.new("RGBA", (self._imgsize, self._imgsize), (0,0,0,0))
        mask_pix = mask.load()

        # Load photo
        layer_pix = self._tile.load() 
        
        # Scan!
        for y in range(self._tile.width-1):
            for x in range(self._tile.height-1):
                p = layer_pix[x,y]
                if p[3] < 255: # <- Check for empty or non-complete alpha
                    mask_pix[x,y] = (0,0,0,255)
        # We do not apply any blur or other effects here - we only want the
        # exact pixel positions.

        #mask.save( mstr_datafolder + "_cache/" + str(self._lat) + "-" + str(self._ty) + "_" + str(self._lng) + "-" + str(self._tx) + "_tile-completion.png" )
        mstr_msg("photogen", "Generated and saved empty space mask")
        return mask


    # Construct a folder name for latitude and longitude
    def latlng_folder(self, numbers):
        fstr = ""
        if numbers[0] >= 0: fstr = "+"
        if numbers[0] <  0: fstr = "-"
        if abs(numbers[0]) < 10: fstr = fstr + "0" + str(numbers[0])
        if abs(numbers[0]) >= 10 and numbers[0] <= 90: fstr = fstr + str(numbers[0])

        if numbers[1] >= 0: fstr = fstr + "+"
        if numbers[1] <  0: fstr = fstr + "-"
        if abs(numbers[1]) < 10: fstr = fstr + "00" + str(numbers[1])
        if abs(numbers[1]) >= 10 and numbers[0] <= 99: fstr = fstr + "0" + str(numbers[1])
        if abs(numbers[1]) >= 100 : fstr = fstr + str(numbers[1])

        return fstr


    # Find the next "by-ten" numbers for the current latitude and longitude
    def find_earthnavdata_number(self):
        earthnavdata = []
        lat = abs(int(self._lat / 10) * 10)
        lng = abs(int(self._lng / 10) * 10)
        earthnavdata.append(lat)
        earthnavdata.append(lng)
        return earthnavdata