From 17640764821d5f272511af531884a84eb3438bbe Mon Sep 17 00:00:00 2001 From: Fabian Gonzalez Date: Sun, 14 Jan 2024 22:26:12 -0500 Subject: [PATCH] More linting - 95% done --- ui/listicles/interface.py | 22 ++++++++++++++--- ui/listicles/utils.py | 52 +++++++++++++++++++++++++++++++++------ ui/music/utils.py | 9 +++---- utils/progress.py | 9 ++++--- utils/visualizer.py | 24 ++++++++++++++++-- 5 files changed, 93 insertions(+), 23 deletions(-) diff --git a/ui/listicles/interface.py b/ui/listicles/interface.py index b27d214..664cd72 100644 --- a/ui/listicles/interface.py +++ b/ui/listicles/interface.py @@ -1,12 +1,18 @@ +""" +The interface for the Listicles section of the web app. +""" +import json import gradio as gr import processing.image as image_processing -import json import ui.listicles.utils as listicle_utils import ui.components.openai as openai_components import utils.gradio as gru -def render_listicles_section(): +def render_listicles_section() -> None: + """ + Renders the Listicles section of the web app. + """ gru.render_tool_description("Create images in the style of those 'Your birth month is your ___' TikToks.") with gr.Tab("Generate Artifacts"): send_artifacts_to_batch_button, listicle_image_output, listicle_json_output = render_generate_section() @@ -20,7 +26,11 @@ def render_listicles_section(): ) -def render_batch_section(): +def render_batch_section() -> (gr.File, gr.Code): + """ + Renders the Batch Image Processing section of the web app. + :return: The input images and input json components. + """ with gr.Column(): gr.Markdown("# Input") with gr.Row(equal_height=False): @@ -106,7 +116,11 @@ def set_json(json_file): return input_batch_images, input_batch_json -def render_generate_section(): +def render_generate_section() -> (gr.Button, gr.Gallery, gr.Code): + """ + Renders the Generate Artifacts section of the web app. + :return: The send artifacts to batch button, the listicle image output gallery, and the listicle json output. + """ api_key, api_text_model, api_image_model = openai_components.render_openai_setup() with gr.Row(equal_height=False): with gr.Group(): diff --git a/ui/listicles/utils.py b/ui/listicles/utils.py index d43eeb5..b4c75b3 100644 --- a/ui/listicles/utils.py +++ b/ui/listicles/utils.py @@ -1,7 +1,11 @@ -import gradio as gr +""" +This file contains the functions that are used by the Gradio UI to generate listicles. +""" +import os import json +from typing import Optional +import gradio as gr import processing.image as image_processing -import os from utils import font_manager, image as image_utils import api.chatgpt as chatgpt_api @@ -93,7 +97,12 @@ def process(image_files, json_data, return images -def validate_json(json_file): +def validate_json(json_file: str) -> None: + """ + Validates the JSON file to make sure it has the required fields. + :param json_file: The JSON file to validate. + :return: None + """ if not json_file or len(json_file) == 0: gr.Warning("No JSON in the code block.") return @@ -117,7 +126,13 @@ def validate_json(json_file): gr.Info("JSON is valid!") -def send_artifacts_to_batch(listicle_images, json_data): +def send_artifacts_to_batch(listicle_images: gr.data_classes.RootModel, json_data: str) -> (list, str): + """ + Sends the artifacts to the batch processing section. + :param listicle_images: The list of images to send. This is a Gradio Gallery. + :param json_data: The JSON data to send. + :return: The list of images and the JSON data sent. + """ if not listicle_images or len(listicle_images.root) == 0: gr.Warning("No images to send.") return @@ -130,10 +145,17 @@ def send_artifacts_to_batch(listicle_images, json_data): return listicle_images, json_data -def save_artifacts(listicle_images, image_type, json_data): +def save_artifacts(listicle_images: gr.data_classes.RootModel, image_type: gr.Dropdown, json_data: str) -> None: + """ + Saves the artifacts to disk. + :param listicle_images: The list of images to save. This is a Gradio Gallery. + :param image_type: The type of image to save. + :param json_data: The JSON data to save. + :return: None + """ if not json_data or len(json_data) == 0: gr.Warning("No JSON data to save.") - return + return None # Save the images save_dir = image_processing.save_images_to_disk(listicle_images, image_type) @@ -148,8 +170,22 @@ def save_artifacts(listicle_images, image_type, json_data): gr.Info(f"Saved generated artifacts to {save_dir}.") -def generate_listicle(api_key, api_text_model, api_image_model, number_of_items, topic, association, - rating_type, details="", generate_artifacts=False): +def generate_listicle(api_key: str, api_text_model: str, api_image_model: str, number_of_items: int, topic: str, + association: str, rating_type: str, details: str = "", generate_artifacts: bool = False) \ + -> (Optional[str], Optional[str], Optional[list[str]]): + """ + Generates a listicle using the OpenAI API. + :param api_key: The OpenAI API key to use. + :param api_text_model: The OpenAI API text model to use (e.g. 'gpt-4'). + :param api_image_model: The OpenAI API image model to use (e.g. 'dall-e-3'). + :param number_of_items: The number of items to generate. + :param topic: The topic of the listicle. + :param association: What each item is associated with. + :param rating_type: What the rating represents. + :param details: Additional details about the listicle you want to generate. + :param generate_artifacts: Whether to generate artifacts (images and JSON) for the listicle. + :return: The listicle content, the listicle JSON, and the listicle images. + """ openai = chatgpt_api.get_openai_client(api_key) if openai is None: gr.Warning("No OpenAI client. Cannot generate listicle.") diff --git a/ui/music/utils.py b/ui/music/utils.py index 7caebca..eb6dd98 100644 --- a/ui/music/utils.py +++ b/ui/music/utils.py @@ -5,18 +5,17 @@ import subprocess import re import time -import cv2 +import tempfile from typing import List, Dict, Optional +import cv2 from moviepy.editor import AudioFileClip +import numpy as np +import librosa from utils import font_manager import utils.image as image_utils -import numpy as np -import tempfile from api import chatgpt as chatgpt_api from processing import image as image_processing -import librosa from utils import progress, visualizer -import cProfile def analyze_audio(audio_path: str, target_fps: int) -> (List[Dict[float, float]], np.ndarray): diff --git a/utils/progress.py b/utils/progress.py index c5a14b5..52569ba 100644 --- a/utils/progress.py +++ b/utils/progress.py @@ -6,9 +6,11 @@ def print_progress_bar(current_iteration: int, total_iterations: int, bar_length start_time: Optional[float] = None, end: str = ''): progress_percentage = (current_iteration / total_iterations) * 100 completed_length = int(bar_length * current_iteration // total_iterations) - bar = '█' * completed_length + '░' * (bar_length - completed_length) + progress_bar = '█' * completed_length + '░' * (bar_length - completed_length) elapsed_time = None + estimated_remaining_time = None + iterations_per_sec = None if start_time is not None: elapsed_time = time.time() - start_time if current_iteration > 0: @@ -19,7 +21,6 @@ def print_progress_bar(current_iteration: int, total_iterations: int, bar_length estimated_remaining_time = None time_string = '' - if estimated_remaining_time is not None: + if estimated_remaining_time is not None and iterations_per_sec is not None: time_string = f'[{elapsed_time:.2f}s/{estimated_remaining_time:.2f}s, {iterations_per_sec:.2f}it/s]' - print(f'\r{progress_percentage:3.0f}%|{bar}| {current_iteration}/{total_iterations} {time_string}', end=end, flush=True) - + print(f'\r{progress_percentage:3.0f}%|{progress_bar}| {current_iteration}/{total_iterations} {time_string}', end=end, flush=True) diff --git a/utils/visualizer.py b/utils/visualizer.py index 8739a70..e32e7bf 100644 --- a/utils/visualizer.py +++ b/utils/visualizer.py @@ -1,8 +1,16 @@ +""" +This module defines the Visualizer class, which is used to draw the visualizer on the canvas. +""" +from typing import Dict, Optional import numpy as np import cv2 class Visualizer: + """ + This class is used to draw the visualizer on the canvas. + Will be replaced with a more general solution in the future to allow for more customization. + """ def __init__(self, base_size, max_size, color, dot_count, width, height): self.base_size = base_size self.max_size = max_size @@ -13,7 +21,11 @@ def __init__(self, base_size, max_size, color, dot_count, width, height): self.cached_dot_positions = None self.cached_resized_drawing = {} - def initialize_static_values(self): + def initialize_static_values(self: "Visualizer") -> None: + """ + Initializes static values for the visualizer. + :return: None. + """ # Calculate and store dot positions x_positions = (self.width / self.dot_count[0]) * np.arange(self.dot_count[0]) + ( self.width / self.dot_count[0] / 2) @@ -23,7 +35,15 @@ def initialize_static_values(self): self.cached_dot_positions = [(grid_x[y, x], grid_y[y, x]) for x in range(self.dot_count[0]) for y in range(self.dot_count[1])] - def draw_visualizer(self, canvas, frequency_data, custom_drawing=None): + def draw_visualizer(self: "Visualizer", canvas: np.ndarray, frequency_data: Dict[float, float], + custom_drawing: Optional[np.ndarray] = None) -> None: + """ + Draws the visualizer on the canvas (a single frame). + :param canvas: The canvas to draw on. + :param frequency_data: The frequency data to use for drawing which correlates to the loudness + frequency. + :param custom_drawing: A custom drawing to use instead of the default circle. + :return: None. + """ # Calculate and store dot positions dot_count_x = self.dot_count[0] dot_count_y = self.dot_count[1]