my_model/dataset/dataset_processor.py

import json
from collections import Counter
import contractions
import csv
import pandas as pd
from typing import Tuple, List, Optional
from my_model.config import dataset_config as config

class OKVQADatasetProcessor:
    """
    Processes the OKVQA dataset by loading, processing, and merging question and annotation data.

    Attributes:
        questions_file_path (str): Path to the questions JSON file.
        annotations_file_path (str): Path to the annotations JSON file.
        questions (List[dict]): Extracted list of question entries from the JSON file.
        annotations (List[dict]): Extracted list of annotation entries from the JSON file.
        df_questions (DataFrame): DataFrame holding the questions.
        df_answers (DataFrame): DataFrame holding the annotations.
        merged_df (Optional[DataFrame]): DataFrame resulting from merging questions and answers, initialized as None.
    """

    def __init__(self, questions_file_path: str, annotations_file_path: str) -> None:
        """
        Initializes the dataset processor with file paths and loads the data into DataFrames.

        Parameters:
            questions_file_path (str): The file path for the questions JSON file.
            annotations_file_path (str): The file path for the annotations JSON file.
        """

        self.questions_file_path = questions_file_path
        self.annotations_file_path = annotations_file_path
        self.questions, self.annotations = self.load_data_files()
        self.df_questions = pd.DataFrame(self.questions)
        self.df_answers = pd.DataFrame(self.annotations)
        self.merged_df = None

    
    def load_data_files(self) -> Tuple[List[dict], List[dict]]:
        """
        Loads the question and annotation data from JSON files.

        Returns:
            Tuple[List[dict], List[dict]]: A tuple containing lists of questions and annotations.
        """
        with open(self.questions_file_path, 'r') as file:
            data = json.load(file)
            questions = data['questions']

        with open(self.annotations_file_path, 'r') as file:
            data = json.load(file)
            annotations = data['annotations']

        return questions, annotations

    
    @staticmethod
    def find_most_frequent(my_list: List[str]) -> Optional[str]:
        """
        Determines the most frequent item in a list.

        Parameters:
            my_list (List[str]): The list from which to find the most frequent item.

        Returns:
            Optional[str]: The most frequent item or None if the list is empty.
        """
        if not my_list:
            return None
        counter = Counter(my_list)
        most_common = counter.most_common(1)
        return most_common[0][0]

    
    def merge_data(self) -> None:
        """
        Merges the question and answer DataFrames on a common key.

        This method sets the 'merged_df' attribute to the resulting DataFrame after merging
        'df_questions' and 'df_answers' on the 'question_id' field, which is assumed to be
        present in both DataFrames.
        """

        self.merged_df = pd.merge(self.df_questions, self.df_answers, on=['question_id', 'image_id'])

    def join_words_with_hyphen(self, sentence):
        
        return '-'.join(sentence.split())

        
    def process_answers(self) -> None:
        """
        Processes answers from merged DataFrame by extracting and identifying the most frequent answers.
        """
        if self.merged_df is not None:
            self.merged_df['raw_answers'] = self.merged_df['answers'].apply(lambda x: [ans['raw_answer'] for ans in x])
            self.merged_df['processed_answers'] = self.merged_df['answers'].apply(
                lambda x: [ans['answer'] for ans in x])
            self.merged_df['most_frequent_raw_answer'] = self.merged_df['raw_answers'].apply(self.find_most_frequent)
            self.merged_df['most_frequent_processed_answer'] = self.merged_df['processed_answers'].apply(
                self.find_most_frequent)
            self.merged_df.drop(columns=['answers'], inplace=True)
        else:
            print("DataFrames have not been merged yet.")

        # Apply the function to the 'most_frequent_processed_answer' column
        self.merged_df['single_word_answers'] = self.merged_df['most_frequent_processed_answer'].apply(
            self.join_words_with_hyphen)

    
    def get_processed_data(self) -> Optional[pd.DataFrame]:
        """
        Retrieves the processed DataFrame.

        Returns:
            Optional[pd.DataFrame]: The processed DataFrame or None if it is not available.
        """

        if self.merged_df is not None:
            return self.merged_df
        else:
            print("DataFrame is empty or not processed yet.")
            return None

    
    def save_to_csv(self, df: pd.DataFrame, saved_file_name: Optional[str]) -> None:
        """
        Saves the DataFrame to a CSV file.

        Parameters:
            df (pd.DataFrame): The DataFrame to save.
            saved_file_name (Optional[str]): The target file name or path.
        """

        if saved_file_name is not None:
            if ".csv" not in saved_file_name:
                df.to_csv(os.path.join(saved_file_name, ".csv"), index=None)
            else:
                df.to_csv(saved_file_name, index=None)
        else:
            df.to_csv("data.csv", index=None)

    
    def display_dataframe(self) -> None:
        """
        Displays the processed DataFrame.
        """
        if self.merged_df is not None:
            print(self.merged_df)
        else:
            print("DataFrame is empty.")



def process_okvqa_dataset(questions_file_path: str, annotations_file_path: str, save_to_csv: bool = False, 
                          saved_file_name: Optional[str] = None) -> Optional[pd.DataFrame]:
    """
    Orchestrates the processing of the OK-VQA dataset using specified JSON file paths for questions and annotations.

    Parameters:
        questions_file_path (str): Path to the questions JSON file.
        annotations_file_path (str): Path to the annotations JSON file.
        save_to_csv (bool): Flag to determine if the processed data should be saved to CSV.
        saved_file_name (Optional[str]): Filename or path to save the CSV file. If None, defaults to 'data.csv'.

    Returns:
        Optional[pd.DataFrame]: The processed DataFrame containing merged and processed VQA data or None if empty.
    """
    # Initialize the dataset processor
    processor = OKVQADatasetProcessor(questions_file_path, annotations_file_path)

    # Merge question and answer data and process answers
    processor.merge_data()
    processor.process_answers()

    # Retrieve the processed DataFrame
    processed_data = processor.get_processed_data()

    # Optionally save the processed DataFrame to a CSV file
    if save_to_csv and processed_data is not None:
        processor.save_to_csv(processed_data, saved_file_name)

    return processed_data