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.")