JabberCard/jabbercard/utilities/graphics.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

from jabbercard.config import Cache
import glob
import os
import qrcode
import random
import xml.etree.ElementTree as ET

try:
    import cv2
except:
    print('OpenCV (cv2) is required for dynamic background.')

try:
    import numpy
except:
    print('NumPy (numpy) is required for dynamic background.')

class Graphics:

    def handle_photo(jid_bare, jid_vcard, link_href):
        filename = filepath = filetype = mimetype = selection = None
        directory_cache = Cache.get_directory()
        filecirca = os.path.join(directory_cache, 'photo', jid_bare, '.*')
        filepath_guess = glob.glob(filecirca)
        if filepath_guess:
            filepath = filepath_guess[0]
            filetype = filepath.split('.').pop()
            filename = '{}.{}'.format(jid_bare, filetype)
        elif jid_vcard:
            if jid_vcard['type']:
                mimetype = jid_vcard['type']
                if mimetype:
                    filetype = mimetype.split('/')[1]
                    if filetype == 'svg+xml': filetype = 'svg'
                    filename = '{}.{}'.format(jid_bare, filetype)
                    filepath = os.path.join(directory_cache, 'photo', filename)
                    #img.save(filename)

                    # Write the decoded bytes to a file
                    if 'bin' in jid_vcard:
                        with open(filepath, 'wb') as file:
                            file.write(jid_vcard['bin'])

        if not filepath or not os.path.exists(filepath) or os.path.getsize(filepath) == 0:
            filename = 'default.svg'
        elif filetype == 'svg':
            selection = Graphics.extract_colours_from_vector(filepath)
        else:
            selection = Graphics.extract_colours_from_raster(filepath)

        # QR code
        filepath_qrcode = os.path.join(directory_cache, 'qr', jid_bare, '.png')
        if not os.path.exists(filepath_qrcode) or os.path.getsize(filepath_qrcode) == 0:
            Graphics.generate_qr_code_graphics_from_string(link_href, jid_bare)

        return filename, filepath, filetype, selection


    def extract_colours_from_raster(filepath):
        try:
            img = cv2.imread(filepath)
            #thresholded = cv2.inRange(img, (50, 100, 200), (50, 100, 200))
            thresholded = cv2.inRange(img, (90, 90, 90), (190, 190, 190))
            #thresholded = cv2.bitwise_not(thresholded)
            #thresholded = cv2.inRange(img, (0, 0, 0), (0, 0, 0))
            #res = img + cv2.cvtColor(thresholded, cv2.COLOR_GRAY2BGR)

            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            #result = numpy.clip(img, 90, 190)
            #result = numpy.clip(img, 50, 200)
            #result = numpy.clip(img, 100, 150)
            result = numpy.clip(img, 100, 200)
            res = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)

            """
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            mask = numpy.all(numpy.logical_and(img >= 90, img <= 190), axis=2)
            result = numpy.where(mask[...,None], img, 255)
            res = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
            """

            """
            # Thresholding for black:
            lower_black = numpy.array([0, 0, 0])
            upper_black = numpy.array([50, 50, 50]) # Adjust this value for the black range 
            black_mask = cv2.inRange(img, lower_black, upper_black)

            # Thresholding for white:
            lower_white = numpy.array([250, 250, 250])
            upper_white = numpy.array([255, 255, 255])
            white_mask = cv2.inRange(img, lower_white, upper_white)

            # Combine the masks
            combined_mask = cv2.bitwise_or(black_mask, white_mask)

            # Invert the combined mask
            inverted_mask = cv2.bitwise_not(combined_mask)

            # Apply the mask to the original image
            res = cv2.bitwise_and(img, img, mask=inverted_mask)
            """

            selection = []

            ix_1st = random.randint(1, len(res)-1)
            res_ix_1st = res[ix_1st]
            ix_ix_1st = random.randint(1, len(res_ix_1st)-1)
            res_ix_ix_1st = res_ix_1st[ix_ix_1st]
            selection.append(numpy.array(res_ix_ix_1st).tolist())

            ix_2nd = random.randint(1, len(res)-1)
            res_ix_2nd = res[ix_2nd]
            ix_ix_2nd = random.randint(1, len(res_ix_2nd)-1)
            res_ix_ix_2nd = res_ix_2nd[ix_ix_2nd]
            selection.append(numpy.array(res_ix_ix_2nd).tolist())
            print(selection)

        except Exception as e:
            selection = None
            exception = str(e)
            print(exception)

        return selection

    def extract_colours_from_vector(filepath):
        # Parse the SVG file
        tree = ET.parse(filepath)
        root = tree.getroot()

        # Set to store unique colours
        colours_hex = set()
        colours_rgb = []

        # SVG namespace
        namespace = {'svg': 'http://www.w3.org/2000/svg'}

        # Find all possible elements
        for elem in root.findall('.//svg:circle', namespace) + \
                    root.findall('.//svg:ellipse', namespace) + \
                    root.findall('.//svg:line', namespace) + \
                    root.findall('.//svg:path', namespace) + \
                    root.findall('.//svg:polygon', namespace) + \
                    root.findall('.//svg:rect', namespace) + \
                    root.findall('.//svg:text', namespace):

            fill = elem.get('fill')
            stroke = elem.get('stroke')

            # Add colours to the set if they are not None or 'none'
            if fill and fill.startswith('#') and len(fill) > 4 and fill.lower() != 'none':
                colours_hex.add(fill)
            if stroke and stroke.startswith('#') and len(stroke) > 4 and stroke.lower() != 'none':
                colours_hex.add(stroke)

        for colour in colours_hex:
            hex = colour.lstrip('#')
            rgb = list(int(hex[i:i+2], 16) for i in (0, 2, 4))
            rgb.reverse()
            colours_rgb.append(rgb)

        selection = []
        if len(colours_rgb) > 1:
            for i in range(2):
                ix = random.randint(0, len(colours_rgb)-1)
                selection.append(colours_rgb[ix])
                del colours_rgb[ix]
        elif len(colours_rgb) == 1:
            selection = [colours_rgb[0], colours_rgb[0]]

        return selection

    def generate_qr_code_graphics_from_string(text, jid_bare):
        #qrcode_graphics = qrcode.make(text)
        qr = qrcode.QRCode(border=2, box_size=10)
        qr.add_data(text)
        qrcode_graphics = qr.make_image(fill_color='#333', back_color='#f2f2f2')
        directory_cache = Cache.get_directory()
        filename = os.path.join(directory_cache, 'qr', jid_bare + '.png')
        qrcode_graphics.save(filename)