Update screenshot.py

screenshot.py

@@ -1,387 +1,56 @@
-import logging
-import math
-from pathlib import Path
-import re
-from time import time
-from typing import List, Optional, Tuple
-
-from PIL import Image, ImageDraw, ImageFont
-from PIL.ImageFont import FreeTypeFont
-
-from PIL import ImageFilter
-logging.basicConfig(level="DEBUG")
-logger = logging.getLogger(__name__)
-
-BG_COMP = "bg_composant"
-PROMPT_VAR = "prompt_variable"
-GENERATION_VAR = "generation_var"
-BOX_COMP = "box_comp"
-
-
-def wrap_text(font: FreeTypeFont, text: str, max_width: int, direction: str = "ltr") -> str:
-    """
-    Wraps the text at the given width.
-    :param font: Font to use.
-    :param text: Text to fit.
-    :param max_width: Maximum width of the final text, in pixels.
-    :param max_height: Maximum height height of the final text, in pixels.
-    :param spacing: The number of pixels between lines.
-    :param direction: Direction of the text. It can be 'rtl' (right to
-                      left), 'ltr' (left to right) or 'ttb' (top to bottom).
-                      Requires libraqm.
-    :return: The wrapped text.
-    """
-    words = text.split()
-
-    lines: list[str] = [""]
-    curr_line_width = 0
-
-    for word in words:
-        if curr_line_width == 0:
-            word_width = font.getlength(word, direction)
-
-            lines[-1] = word
-            curr_line_width = word_width
-        else:
-            new_line_width = font.getlength(f"{lines[-1]} {word}", direction)
-
-            if new_line_width > max_width:
-                # Word is too long to fit on the current line
-                word_width = font.getlength(word, direction)
-
-                # Put the word on the next line
-                lines.append(word)
-                curr_line_width = word_width
-            else:
-                # Put the word on the current line
-                lines[-1] = f"{lines[-1]} {word}"
-                curr_line_width = new_line_width
-
-    return "\n".join(lines)
-
-
-# tr.wrap(my_str, width=30)
-
-
-def _fit_paragraph(
-    paragraph,
-    font: FreeTypeFont,
-    max_width: int,
-    max_height: int,
-    line_height,
-    num_lines: List[str],
-    spacing: int = 4,
-    direction: str = "ltr",
-    
-):
-    paragraph_lines: list[str] = [""]
-    curr_line_width = 0
-    # This is a very bad splitter for Chinese
-    words = list(paragraph) if re.search(u'[\u4e00-\u9fff]',paragraph) else paragraph.split(" ")
-
-    for word in words:
-        if curr_line_width == 0:
-            word_width = font.getlength(word, direction)
-
-            if word_width > max_width:
-                # Word is longer than max_width
-                return None
-
-            paragraph_lines[-1] = word
-            curr_line_width = word_width
-        else:
-            new_line_width = font.getlength(f"{paragraph_lines[-1]} {word}", direction)
-
-            if new_line_width > max_width:
-                # Word is too long to fit on the current line
-                word_width = font.getlength(word, direction)
-                new_num_lines = num_lines + len(paragraph_lines) + 1
-                new_text_height = (new_num_lines * line_height) + (new_num_lines * spacing)
-
-                if word_width > max_width or new_text_height > max_height:
-                    # Word is longer than max_width, and
-                    # adding a new line would make the text too tall
-                    return None
-
-                # Put the word on the next line
-                paragraph_lines.append(word)
-                curr_line_width = word_width
-            else:
-                # Put the word on the current line
-                paragraph_lines[-1] = f"{paragraph_lines[-1]} {word}"
-                curr_line_width = new_line_width
-    return paragraph_lines
-
-def try_fit_text(
-    font: FreeTypeFont,
-    prompt: str,
-    generation: str,
-    max_width: int,
-    max_height: int,
-    spacing: int = 4,
-    direction: str = "ltr",
-) -> Optional[str]:
-    """
-    Attempts to wrap the text into a rectangle.
-    Tries to fit the text into a box using the given font at decreasing sizes,
-    based on ``scale_factor``. Makes ``max_iterations`` attempts.
-    :param font: Font to use.
-    :param text: Text to fit.
-    :param max_width: Maximum width of the final text, in pixels.
-    :param max_height: Maximum height height of the final text, in pixels.
-    :param spacing: The number of pixels between lines.
-    :param direction: Direction of the text. It can be 'rtl' (right to
-                      left), 'ltr' (left to right) or 'ttb' (top to bottom).
-                      Requires libraqm.
-    :return: If able to fit the text, the wrapped text. Otherwise, ``None``.
-    """
-
-    line_height = font.size
-
-    if line_height > max_height:
-        # The line height is already too big
-        return None
-
-    prompt_lines: list[str] = []
-    paragraphs = prompt.split("\n")
-    for paragraph in paragraphs:
-        paragraph_lines = _fit_paragraph(
-            paragraph=paragraph,
-            font=font,
-            max_width=max_width,
-            max_height=max_height,
-            line_height=line_height,
-            spacing=spacing,
-            num_lines=len(prompt_lines),
-            direction=direction,
-        )
-        if paragraph_lines is None:
-            return None
-        prompt_lines.extend(paragraph_lines)
-    generation_lines: list[str] = []
-    paragraphs = f"{prompt_lines[-1]}{generation}".split("\n")
-    for paragraph in paragraphs:
-        paragraph_lines = _fit_paragraph(
-            paragraph=paragraph,
-            font=font,
-            max_width=max_width,
-            max_height=max_height,
-            line_height=line_height,
-            spacing=spacing,
-            num_lines=len(prompt_lines) + len(generation_lines),
-            direction=direction,
-        )
-        if paragraph_lines is None:
-            return None
-        generation_lines.extend(paragraph_lines)
-    generation_lines[0] = generation_lines[0][len(prompt_lines[-1]):]
-    return "\n".join(prompt_lines), "\n".join(generation_lines)
-
-
-# pylint: disable=too-many-arguments
-def fit_text(
-    font,
-    prompt: str,
-    generation: str,
-    max_width: int,
-    max_height: int,
-    spacing: int = 4,
-    scale_factor: float = 0.8,
-    max_iterations: int = 5,
-    direction: str = "ltr",
-) -> Tuple[FreeTypeFont, str]:
-    """
-    Automatically determines text wrapping and appropriate font size.
-    Tries to fit the text into a box using the given font at decreasing sizes,
-    based on ``scale_factor``. Makes ``max_iterations`` attempts.
-    If unable to find an appropriate font size within ``max_iterations``
-    attempts, wraps the text at the last attempted size.
-    :param font: Font to use.
-    :param text: Text to fit.
-    :param max_width: Maximum width of the final text, in pixels.
-    :param max_height: Maximum height height of the final text, in pixels.
-    :param spacing: The number of pixels between lines.
-    :param scale_factor:
-    :param max_iterations: Maximum number of attempts to try to fit the text.
-    :param direction: Direction of the text. It can be 'rtl' (right to
-                      left), 'ltr' (left to right) or 'ttb' (top to bottom).
-                      Requires libraqm.
-    :return: The font at the appropriate size and the wrapped text.
-    """
-    initial_font_size = font.size
-
-    # logger.debug('Trying to fit text "%s"', text)
-
-    for i in range(max_iterations):
-        trial_font_size = int(initial_font_size * pow(scale_factor, i))
-        trial_font = font.font_variant(size=trial_font_size)
-
-        logger.debug("Trying font size %i", trial_font_size)
-
-        wrapped = try_fit_text(trial_font, prompt, generation, max_width, max_height, spacing, direction)
-
-        if wrapped is not None:
-            logger.debug("Successfully fit text")
-            return (trial_font, wrapped)
-
-    # Give up and wrap the text at the last size
-    logger.debug("Gave up trying to fit text; just wrapping text")
-    wrapped = wrap_text(trial_font, prompt, max_width, direction) + wrap_text(
-        trial_font, generation, max_width, direction
-    )
-
-    return (trial_font, wrapped)
-
-
-
-def main(
-    prompt,
-    generation,
-    width,
-    height,
-    assets_path,
-    font_path,
-    colors,
-    frame_to_box_margin,
-    text_to_text_box_margin,
-    init_font_size,
-    right_align = False
-):
-    # prompt_color = "#ffffff"
-    # text_color = "#FF57A0"
-    # text_box = ((500, 500, 2300, 1800))
-    # margin = 50
-    # input_color = "#CDD2E3"
-    right_align_params = {"direction":'rtl',"align":'right',"features":'rtla'} if right_align else {}
-    text_box_margin_r = frame_to_box_margin
-    text_box_margin_t = frame_to_box_margin
-    text_box_margin_l = frame_to_box_margin
-    text_box_margin_b = int(height / 4.5)
-
-    text_box = (
-        text_box_margin_l,
-        text_box_margin_t,
-        width - text_box_margin_r,
-        height - text_box_margin_b,
-    )
-
-    background = Image.new("RGB", (width, height), color=colors[BG_COMP])
-    # Get assets
-    assets_path = Path(assets_path)
-    flower_1 = Image.open(assets_path / "image_1.png").copy()
-    flower_2 = Image.open(assets_path / "image_2.png").copy()
-    shadow = Image.open(assets_path / "image_3.png").copy()
-    bloom_logo = Image.open(assets_path / "image_4.png").copy()
-    input_info = Image.open(assets_path / "image_7.png").copy()
-    output_info = Image.open(assets_path / "image_6.png").copy()
-    details = Image.open(assets_path / "image_5.png").copy()
-
-    flower_1_offsets = (int(width - flower_1.width * 2 / 3), int(-flower_1.height / 3))
-    background.paste(flower_1, flower_1_offsets, flower_1)
-
-    flower_2_offsets = (
-        -int(flower_2.width * 2 / 5),
-        int(height / 2 - flower_2.height / 2),
-    )
-    background.paste(flower_2, flower_2_offsets, flower_2)
-
-    bloom_offsets = (
-        frame_to_box_margin,
-        int(height - bloom_logo.height - frame_to_box_margin),
-    )
-    background.paste(bloom_logo, bloom_offsets, bloom_logo)
-
-    input_info_offsets = (
-        width - details.width - text_to_text_box_margin - input_info.width - output_info.width ,
-        int(height - details.height - frame_to_box_margin),
-    )
-    background.paste(input_info, input_info_offsets, input_info)
-
-    output_info_offsets =(
-        width - details.width - text_to_text_box_margin - input_info.width - output_info.width ,
-        int(height - details.height - frame_to_box_margin + input_info.height + text_to_text_box_margin),
-    )
-    background.paste(output_info, output_info_offsets, output_info)
-
-    details_offsets = (
-        width - frame_to_box_margin - details.width ,
-        int(height - details.height - frame_to_box_margin),
-    )
-    background.paste(details, details_offsets, details)
-    box_margin = (
-        text_box[0] + text_to_text_box_margin,
-        text_box[1] + text_to_text_box_margin,
-        text_box[2] - text_to_text_box_margin,
-        text_box[3] - text_to_text_box_margin,
-    )
-
-    # text_box = ImageText(box_margin)
-    drawing = ImageDraw.Draw(background, "RGBA")
-
-    # Text box for main text
-    input_color = colors[BOX_COMP][1:]
-    input_color = tuple(int(input_color[i : i + 2], 16) for i in (0, 2, 4)) + (
-        int(255 * 0.99),
-    )  # RGB + A
-    drawing.rounded_rectangle(text_box, outline="#000", fill=input_color, radius=47.9)
-
-    # Adapt text size to box
-    # font, (prompt_a, generation_a) = adapt_text_to_ratio(box_margin, prompt, generation)
-    # generation_a = adapt_text_to_ratio(box_margin, generation, prompt)
-    # font = adapt_font_to_text(box_margin, prompt_a + generation_a, font_path=font_path)
-
-    init_font = ImageFont.truetype(font_path, size=init_font_size)
-    font, (prompt_a, generation_a) = fit_text(
-        prompt=prompt,
-        generation=generation,
-        font=init_font,
-        max_height=box_margin[3] - box_margin[1],
-        max_width=box_margin[2] - box_margin[0],
-        max_iterations=50,
-        scale_factor=0.95,
-        direction="rtl" if right_align else "ltr"
-    )
-    # Prompt, main, then the last line
-    prompt_s = prompt_a.split("\n")
-    prompt_main = "\n".join(prompt_s[:-1])
-    prompt_lastline = prompt_s[-1]
-    drawing.multiline_text(
-        (box_margin[0], box_margin[1]), 
-        prompt_main, 
-        colors[PROMPT_VAR], 
-        font,
-        **right_align_params
-    )
-    end_prompt_main = font.getsize_multiline(prompt_main)
-    end_prompt_last = font.getsize_multiline(prompt_lastline)
-    drawing.multiline_text(
-        ((box_margin[2] - end_prompt_last[0]) if right_align else box_margin[0], box_margin[1] + end_prompt_main[1]),
-        prompt_lastline,
-        colors[PROMPT_VAR],
-        font,
-        **right_align_params
-    )
-
-    # Generated text, first line, then the rest
-    generation_split = generation_a.split("\n")
-    generation_firstline = generation_split[0]
-    generation_main = "\n".join(generation_split[1:])
-    drawing.multiline_text(
-        # margin x + length(last line of prompt), margin Y + length(main part of prompt)
-        (box_margin[0], box_margin[1] + end_prompt_main[1]) if right_align else \
-        (box_margin[0] + end_prompt_last[0], box_margin[1] + end_prompt_main[1]),
-        generation_firstline,
-        colors[GENERATION_VAR],
-        font, 
-        **right_align_params
-    )
-    drawing.multiline_text(
-        (box_margin[0], box_margin[1] + end_prompt_main[1] + end_prompt_last[1]),
-        generation_main,
-        colors[GENERATION_VAR],
-        font,
-        **right_align_params
-    )
-
-    final_font_size = font.size
-    return final_font_size, background
+## HTML and JS code to give Gradio HTML
+before_prompt = """
+  <div id = "img_placeholder">
+  </div>
+  <div class="relative" id="capture" align="justify" style="display:none;">
+    <div class="absolute font-semibold" style="left:7%; right:7%; bottom:32%; top:7%; font-size: 8rem; line-height: 1; padding: 1rem; font-family:-apple-system, BlinkMacSystemFont, 'Arial', sans-serif;" id="text_box">
+      <p class="text" style="color:white; white-space:pre-wrap;" dir="auto" id = "prompt">"""
+prompt_to_generation = """</p>
+        <p class="text" style="color:#FE57A0; white-space:pre-wrap;" dir="auto" id="generation">"""
+after_generation = """</p>
+    </div>
+    <img src="https://huggingface.co/spaces/clefourrier/bloom_media_card/raw/main/bg.jpg" class="w-full" />
+  </div>
+"""
+
+js_save = """() => {
+      /*might need to add .root to launch locally */
+      var gradioapp = document.body.getElementsByTagName('gradio-app')[0]; 
+      
+      /* Save image */
+      capture = gradioapp.querySelector('#capture')
+      img_placeholder = gradioapp.querySelector('#img_placeholder')
+      html2canvas(capture, {
+            useCORS: true,
+            onclone: function (clonedDoc) {
+                  clonedDoc.querySelector('#capture').style.display = 'block';
+                  
+                  /*Fits text to box*/
+                  var text_box = clonedDoc.querySelector('#text_box');
+                  var prompt = clonedDoc.querySelector('#prompt');
+                  var generation = clonedDoc.querySelector('#generation');
+                  console.log(text_box, generation, prompt)
+                  cur_font_size = getComputedStyle(text_box).getPropertyValue("font-size")
+                  while( (text_box.clientHeight < text_box.scrollHeight || text_box.clientWidth < text_box.scrollWidth) &  parseInt(cur_font_size) > 10) {
+                    console.log(cur_font_size, text_box.clientHeight, text_box.scrollHeight, text_box.clientWidth, text_box.scrollWidth)
+                    cur_font_size = 0.98 * parseInt(cur_font_size) + "px"
+                    cur_line_height = 1.1 * parseInt(cur_font_size) + "px"
+                    text_box.style.fontSize = cur_font_size
+                    prompt.style.fontSize = cur_font_size
+                    generation.style.fontSize = cur_font_size
+                    text_box.style.lineHeight = cur_line_height
+                    prompt.style.lineHeight = cur_line_height
+                    generation.style.lineHeight = cur_line_height
+                  }
+            }
+        }).then((canvas)=>{
+          img_placeholder.prepend(canvas);
+        })
+}"""
+
+
+js_load_script="""() => {
+                  var script = document.createElement('script');
+                  script.src = "https://cdnjs.cloudflare.com/ajax/libs/html2canvas/1.4.1/html2canvas.min.js";
+                  document.head.appendChild(script);
+}"""