vouchervision/utils_hf.py

import os, json, re
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload
from google.oauth2 import service_account
import base64
from PIL import Image
from PIL import Image
from io import BytesIO

from vouchervision.general_utils import get_cfg_from_full_path


def setup_streamlit_config(dir_home):
    # Define the directory path and filename
    dir_path = os.path.join(dir_home, ".streamlit")
    file_path = os.path.join(dir_path, "config.toml")

    # Check if directory exists, if not create it
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
    
    # Create or modify the file with the provided content
    config_content = f"""
    [theme]
    base = "dark"
    primaryColor = "#00ff00"

    [server]
    enableStaticServing = false
    runOnSave = true
    port = 8524
    maxUploadSize = 5000
    """

    with open(file_path, "w") as f:
        f.write(config_content.strip())



def save_uploaded_file(directory, img_file, image=None):
    if not os.path.exists(directory):
        os.makedirs(directory)
    # Assuming the uploaded file is an image
    if image is None:
        with Image.open(img_file) as image:
            full_path = os.path.join(directory, img_file.name)
            image.save(full_path, "JPEG")
        # Return the full path of the saved image
        return full_path
    else:
        full_path = os.path.join(directory, img_file.name)
        image.save(full_path, "JPEG")
        return full_path
    
def image_to_base64(img):
    buffered = BytesIO()
    img.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode()

def check_prompt_yaml_filename(fname):
    # Check if the filename only contains letters, numbers, underscores, and dashes
    pattern = r'^[\w-]+$'
    
    # The \w matches any alphanumeric character and is equivalent to the character class [a-zA-Z0-9_].
    # The hyphen - is literally matched.

    if re.match(pattern, fname):
        return True
    else:
        return False
    
# Function to upload files to Google Drive
def upload_to_drive(filepath, filename):
    # Parse the service account info from the environment variable
    creds_info = os.environ.get('GDRIVE_API')
    if creds_info:
        creds_info = json.loads(creds_info)
        creds = service_account.Credentials.from_service_account_info(
            creds_info, scopes=["https://www.googleapis.com/auth/drive"]
        )
        service = build('drive', 'v3', credentials=creds)

        # Get the folder ID from the environment variable
        folder_id = os.environ.get('GDRIVE_FOLDER_ID')  # Renamed for clarity

        if folder_id:
            file_metadata = {
                'name': filename,
                'parents': [folder_id]
            }

            # Determine the mimetype based on the file extension
            if filename.endswith('.yaml') or filename.endswith('.yml'):
                mimetype = 'application/x-yaml'
            elif filename.endswith('.zip'):
                mimetype = 'application/zip'
            else:
                # Set a default mimetype if desired or handle the unsupported file type
                print("Unsupported file type")
                return None

            # Upload the file
            media = MediaFileUpload(filepath, mimetype=mimetype)
            file = service.files().create(
                body=file_metadata,
                media_body=media,
                fields='id'
            ).execute()
            print(f"Uploaded file with ID: {file.get('id')}")
    else:
        print("GDRIVE_API environment variable not set.")