logic/data_utils.py

import os
import json
import time
from huggingface_hub import HfApi, create_repo, CommitScheduler
import bcrypt
import shutil
import uuid
import gradio as gr
from PIL import Image
import numpy as np


def load_concepts(path="data/concepts.json"):
    with open(path, encoding='utf-8') as f:
        data = json.load(f)
    sorted_data = dict()
    for country in sorted(data):
        sorted_data[country] = dict()
        for lang in sorted(data[country]):
            sorted_data[country][lang] = data[country][lang]
    return sorted_data


def load_metadata(path="data/metadata.json"):
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    sorted_data = dict()
    for country in sorted(data):
        sorted_data[country] = dict()
        for lang in sorted(data[country]):
            sorted_data[country][lang] = data[country][lang]
    return sorted_data


class CustomHFDatasetSaver:
    def __init__(self, api_token, dataset_name, private=False):
        self.api_token = api_token
        self.dataset_name = dataset_name
        self.private = private
        self.api = HfApi()
        
    def setup(self, data_outputs, local_ds_folder):
        # create repo is not exist
        self.dataset_name = create_repo(
            repo_id=self.dataset_name,
            token=self.api_token,
            private=self.private,
            repo_type="dataset",
            exist_ok=True,
        ).repo_id

        # Create the local data folder if not exist
        self.local_ds_folder = local_ds_folder
        os.makedirs(self.local_ds_folder, exist_ok=True)

        self.data_outputs = data_outputs  # list of components to read values from

        # create scheduler to commit the data to the hub every x minutes
        self.scheduler = CommitScheduler(
            repo_id=self.dataset_name,
            repo_type="dataset",
            folder_path=self.local_ds_folder,
            every=1,
            token=self.api_token,
        )
    
    def validate_data(self, values_dic):
        """
        Validates the data before saving to ensure no required fields are empty.
        Returns (bool, str) tuple where first value indicates if validation passed
        and second value contains error message if validation failed.
        """
        # Remove 'image' from required fields since we handle it separately
        required_fields = ['country', 'language', 'category', 'concept', 'caption']
        
        # Check if image is provided (either uploaded or via URL)
        image = values_dic.get('image')
        image_url = values_dic.get('image_url')
        
        # Check if image exists and is not None
        has_image = image is not None and (isinstance(image, dict) or (hasattr(image, 'shape') and image.shape[0] > 0))
        has_url = image_url is not None and image_url.strip() != ""
        
        if not has_image and not has_url:
            return False, "Either an image or image URL must be provided"
            
        # Check required fields
        for field in required_fields:
            value = values_dic.get(field)
            if value is None or (isinstance(value, str) and value.strip() == ""):
                return False, f"Required field '{field}' cannot be empty"
                
        # Check if image file exists if image path is provided
        if has_image and isinstance(image, dict):
            if not os.path.exists(image.get('path', '')):
                return False, "Image file not found"
                
        return True, ""

    #TODO: add a function to check if the user is logged in
    def is_logged_in(self):
        pass

    #TODO: check if the user is logged in (add a decorator to the save function)
    def save(self, *values):
        # 'values' are the outputs from your data collection components,
        # you can map these to field names as needed
        values_dic = dict(zip(self.data_outputs, values))
        
        # print(f"Values received: {values_dic}")
        
        # Validate data before proceeding
        is_valid, error_msg = self.validate_data(values_dic)
        if not is_valid:
            raise gr.Error(error_msg)
            # raise ValueError(error_msg)
            
        values_dic['password'] = self.hash_password(values_dic['password'])

        # # Process main category and concept
        # main_category = values_dic.get('category', '')
        # main_concept = values_dic.get('concept', '')
        
        # # Process category-specific concept dropdowns
        # additional_concepts_by_category = {}
        # # Extract predefined categories and their corresponding dropdowns from values_dic
        # predefined_categories = sorted(list(values_dic.get('concepts_dict', {})
        #                             .get(values_dic.get('country', 'USA'), {})
        #                             .get(values_dic.get('language', 'English'), {}).keys()))[:5]
        # # Process each category dropdown
        # for i, category in enumerate(predefined_categories):
        #     dropdown_key = f'category{i+1}_concepts'
        #     if dropdown_key in values_dic and values_dic[dropdown_key]:
        #         # Only add non-empty concept selections
        #         if values_dic[dropdown_key]:
        #             additional_concepts_by_category[category] = values_dic[dropdown_key]

        ### TODO: fix saving additional concepts if not displayed in English
        # # Process category-specific concept dropdowns
        # additional_concepts_by_category = {}

        # # Extract the country and language
        # country = values_dic.get('country', 'USA')
        # language = values_dic.get('language', 'English')
        # concepts_dict = values_dic.get('concepts_dict', {})
        # lang2eng_mapping = values_dic.get('country_lang_map', {})

        # # Get the English version of the language for dictionary lookup
        # eng_lang = lang2eng_mapping.get(language, language)

        # # Get the predefined categories in English
        # predefined_categories = sorted(list(concepts_dict.get(country, {}).get(eng_lang, {}).keys()))[:5]

        # # Process each category dropdown
        # for i, category in enumerate(predefined_categories):
        #     dropdown_key = f'category_{i+1}_concepts'
        #     if dropdown_key in values_dic and values_dic[dropdown_key]:
        #         # Only add non-empty concept selections
        #         additional_concepts_by_category[category] = values_dic[dropdown_key]

        current_timestamp = int(time.time() * 1000)
        # Create a unique ID for the sample is not provided
        if not values_dic.get("id"):
            # Missing ID
            country, language, category, concept = values_dic.get("country"), values_dic.get("language"), values_dic.get("category"), values_dic.get("concept")
            values_dic["id"] = f'{country}_{language}_{category}_{concept}_{current_timestamp}'

        #prepare the main directory of the sample
        if values_dic.get("username"):
            sample_dir = os.path.join("logged_in_users", values_dic["country"], values_dic["language"], values_dic["username"], str(current_timestamp))
        else:
            sample_dir = os.path.join("anonymous_users", values_dic["country"], values_dic["language"], str(uuid.uuid4()))

        os.makedirs(os.path.join(self.local_ds_folder, sample_dir), exist_ok=True)

        # Destination path 
        dest_image_path = os.path.join(sample_dir, "image.png")
        # Source path (to be used for copying the file in the with lock block)
        # This is the path of the image file that was uploaded by the user
        # I want to save the values_dic['image'] in the dest_image_path
        # Convert numpy array to PIL Image and save it
        # ===
        # uploaded_image_path = os.path.join(self.local_ds_folder, dest_image_path)
        # img = Image.fromarray(values_dic['image'])
        # img.save(uploaded_image_path)
        full_dest_path = os.path.join(self.local_ds_folder, dest_image_path)
        
        # Handle different image types
        image_data = values_dic['image']
        
        if isinstance(image_data, dict) and 'path' in image_data:
            # New upload case - copy from the uploaded path
            uploaded_image_path = image_data['path']
            with self.scheduler.lock:
                shutil.copy(uploaded_image_path, full_dest_path)
        elif isinstance(image_data, np.ndarray):  # not values_dic.get('excluded', False) and 
            # Exclude case with numpy array - save the array as an image
            with self.scheduler.lock:
                # Convert numpy array to PIL image and save
                img = Image.fromarray(image_data)
                img.save(full_dest_path)
        elif isinstance(image_data, Image.Image):
            # PIL image case
            with self.scheduler.lock:
                image_data.save(full_dest_path)
        
        values_dic['image'] = dest_image_path
        image_file_path_on_hub =  f"https://huggingface.co/datasets/{self.dataset_name}/resolve/main/{dest_image_path}"

        # print(f"Saving sample: {values}")
        # Build the metadata dictionary.
        data_dict = {
            # in case using windows
            "image": values_dic['image'].replace("\\", "/"),
            "image_file": image_file_path_on_hub.replace("\\", "/"),
            # "image": values_dic['image'],
            # "image_file": image_file_path_on_hub,
            "image_url": values_dic['image_url'] or "",
            "caption": values_dic['caption'] or "",
            "country": values_dic['country'] or "",
            "language": values_dic['language'] or "",
            "category": values_dic['category'] or "",
            "concept": values_dic['concept'] or "",
            "category_1_concepts": [""] if values_dic.get('category_1_concepts', [""])==[] else values_dic.get('category_1_concepts', [""]),
            "category_2_concepts": [""] if values_dic.get('category_2_concepts', [""])==[] else values_dic.get('category_2_concepts', [""]),
            "category_3_concepts": [""] if values_dic.get('category_3_concepts', [""])==[] else values_dic.get('category_3_concepts', [""]),
            "category_4_concepts": [""] if values_dic.get('category_4_concepts', [""])==[] else values_dic.get('category_4_concepts', [""]),
            "category_5_concepts": [""] if values_dic.get('category_5_concepts', [""])==[] else values_dic.get('category_5_concepts', [""]),
            "timestamp": current_timestamp,
            "username": values_dic['username'] or "",
            "password": values_dic['password'] or "",
            "id": values_dic['id'],
            "excluded": False if values_dic.get('excluded') is None else bool(values_dic.get('excluded')),
            # "is_blurred": str(values_dic.get('is_blurred'))
        }

        print(f"Data dictionary: {data_dict}")
        # Define a unique filename for the JSON metadata file (stored in self.folder).
        json_filename = f"sample_{current_timestamp}.json"
        json_file_path = os.path.join(self.local_ds_folder, sample_dir, json_filename)

        with self.scheduler.lock:
            # Save the metadata to the sample file in the local dataset folder
            with open(json_file_path, "w", encoding="utf-8") as f:
                json.dump(data_dict, f, indent=2)
        
        print("Data saved successfully")

    def hash_password(self, raw_password):
        """
        Hashes a raw password using bcrypt and returns the hashed password.
            raw_password (str): The plain text password to be hashed.
            str: The hashed password as a string.
        """
        hashed_password = bcrypt.hashpw(raw_password.encode(), bcrypt.gensalt()).decode()
        return hashed_password