{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# **📈 📉 Hybrid model with DONUT and ResNet50**\n## Combining Computer Vision and NLP for Graph Interpretation. ","metadata":{}},{"cell_type":"markdown","source":"# **❗️ Important**\n\n**The models have already been trained and all credits go to their creators, follow the link below to access the notebooks from which they were taken:**\n- **[Tuned Donut](https://www.kaggle.com/code/cody11null/tuned-donut)**\n- **[[Simple Submission]Classification Task - Inference](https://www.kaggle.com/code/bibanh/simple-submission-classification-task-inference)**\n\n**Here we will only intend to join the two models and analyze how it was done, if you want to train your own model and use this code replace the variables:**<br>\n- ***model_donut_dir*** <br> \n- ***model_ResNet50_dir***","metadata":{}},{"cell_type":"markdown","source":"# **📚 About**\n\nThe first component of our hybrid model is the **Donut Model**, an NLP model that specializes in graph interpretation. This model, trained to predict the structure and data of a graph from textual descriptions, is able to provide a first approximation of the type of graph and data series present in the image.\n\nThe second component is an image classification **model based on the ResNet50 architecture**, a state-of-the-art computer vision model. This model is trained to classify the graph type directly from the image, providing an alternative prediction that complements that of the Donut Model.\n\nBy combining the predictions of these two models, we **created a hybrid model** that leverages the strengths of both approaches. The graph data structure is interpreted through the Donut Model, while the graph type is validated and complemented by the ResNet50 image classification.\n\n---","metadata":{}},{"cell_type":"markdown","source":"# **1. Dependencies and code setup**","metadata":{}},{"cell_type":"code","source":"import os\nimport re\nimport cv2 as cv\nimport torch\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport cv2 as cv\nimport random\nimport numpy as np\nimport pandas as pd\nimport torchvision\nimport torch.utils.data\nfrom typing import List\nfrom pathlib import Path\nfrom functools import partial\nimport torchvision.transforms as transforms\nfrom transformers import (\n    DonutProcessor,\n    VisionEncoderDecoderConfig,\n    VisionEncoderDecoderModel,\n)\nfrom torch.utils.data import DataLoader\nfrom datasets import Dataset\nfrom datasets import Image as ds_img\nfrom tqdm.notebook import tqdm\nfrom PIL import Image\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"execution":{"iopub.status.busy":"2023-05-24T15:08:45.257904Z","iopub.execute_input":"2023-05-24T15:08:45.258552Z","iopub.status.idle":"2023-05-24T15:08:57.939847Z","shell.execute_reply.started":"2023-05-24T15:08:45.258518Z","shell.execute_reply":"2023-05-24T15:08:57.938864Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Global Configuration\nclass CFG:\n    # If set to True, performs dataset cleaning for debugging purposes\n    debug_clean = False\n\n    # Batch size for training\n    batch_size = 4\n\n    # Path to the input images\n    image_path = \"/kaggle/input/benetech-making-graphs-accessible/test/images\"\n    \n    # Maximum length for input sequences\n    max_length = 512\n\n    # Directory containing the trained model 1 (DONUT)\n    model_donut_dir = \"/kaggle/input/benetech-donut\"\n    \n    # Directory containing the trained model 2 (ResNet50)\n    model_ResNet50_dir = \"/kaggle/input/benetech-classification-pretrained/Benetech _ResNet50_fold0.pth\"\n\n    # Label mapping to class names\n    label_map = {0: 'dot', 1: 'horizontal_bar', 2: 'vertical_bar', 3: 'line', 4: 'scatter'}\n\n    # Total number of classes\n    num_classes = 5\n\n    # Special tokens to delimit x and y coordinates in input sequences\n    X_START = \"<x_start>\"\n    X_END = \"<x_end>\"\n    Y_START = \"<y_start>\"\n    Y_END = \"<y_end>\"\n\n    # Default data for data series and chart type\n    PLACEHOLDER_DATA_SERIES = \"0;0\"\n    PLACEHOLDER_CHART_TYPE = \"line\"","metadata":{"execution":{"iopub.status.busy":"2023-05-24T15:08:57.941639Z","iopub.execute_input":"2023-05-24T15:08:57.942399Z","iopub.status.idle":"2023-05-24T15:08:57.950646Z","shell.execute_reply.started":"2023-05-24T15:08:57.942365Z","shell.execute_reply":"2023-05-24T15:08:57.949223Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **2. DONUT model** \nThe **DonutDataProcessor** class is defined in this cell. This class is responsible for loading the DONUT model, preparing the data for the model and generating predictions.","metadata":{}},{"cell_type":"code","source":"class DonutDataProcessor:\n    def __init__(self, cfg):\n        \"\"\"\n        Initializes the DonutDataProcessor object.\n\n        Parameters\n        ----------\n        cfg : object\n            Configuration object.\n        \"\"\"\n        self.cfg = cfg\n        self.device = torch.device(\"cuda:0\")\n        self.model = None\n        self.processor = None\n        self.decoder_start_token_id = None\n        self.ds = None\n        self.data_loader = None\n\n    @staticmethod\n    def clean(str_list: List[str]):\n        \"\"\"\n        Cleans a list of stringified numbers by removing unwanted characters and \n        converting to numerical format.\n\n        Parameters\n        ----------\n        str_list : list of str\n            The list of stringified numbers to clean.\n\n        Returns\n        -------\n        list of float\n            A new list where each entry is the cleaned numerical version of \n            the corresponding entry in the input list.\n        \"\"\"\n        new_list = []\n        for temp in str_list:\n            if \".\" not in temp:\n                dtype = int\n            else:\n                dtype = float\n            try:\n                temp = dtype(re.sub(\"\\s\", \"\", temp))\n            except ValueError:\n                temp = re.sub(r\"[^0-9\\.\\-eE]\", \"\", temp)\n\n                if len(temp) == 0:\n                    temp = 0\n                else:\n                    multiple_periods = len(re.findall(r\"\\.\", temp)) > 1\n                    multiple_negative_signs = len(re.findall(r\"\\-\", temp)) > 1\n                    multiple_e = len(re.findall(r\"[eE]\", temp)) > 1\n\n                    if multiple_negative_signs:\n                        temp = \"-\" + re.sub(r\"\\-\", \"\", temp)\n\n                    if multiple_periods:\n                        chunks = temp.split(\".\")\n                        try:\n                            temp = chunks[0] + \".\" + \"\".join(chunks[1:])\n                        except IndexError:\n                            temp = \"\".join(chunks)\n\n                    if multiple_e:\n                        while temp.lower().startswith(\"e\"):\n                            temp = temp[1:]\n\n                        while temp.lower().endswith(\"e\"):\n                            temp = temp[:-1]\n\n                        chunks = temp.split(\"e\")\n                        try:\n                            temp = chunks[0:-1] + \"e\" + \"\".join(chunks[-1])\n                        except IndexError:\n                            temp = \"\".join(chunks)\n                try:\n                    temp = dtype(temp)\n                except ValueError:\n                    temp = 0\n\n            new_list.append(temp)\n\n        return new_list\n\n    def clean_preds(self, x: List[str], y: List[str]):\n        \"\"\"\n        This function cleans the x and y values predicted by Donut.\n\n        Because it is a generative model, it can insert any character in the\n        model's vocabulary into the prediction string. This function primarily removes\n        characters that prevent a number from being cast to a float.\n\n        Parameters\n        ----------\n        x : list of str\n            The x values predicted by Donut.\n        y : list of str\n            The y values predicted by Donut.\n\n        Returns\n        -------\n        list of str, list of str\n            The cleaned x values and cleaned y values.\n        \"\"\"\n        all_x_chars = \"\".join(x)\n        all_y_chars = \"\".join(y)\n\n        frac_num_x = len(re.sub(r\"[^\\d]\", \"\", all_x_chars)) / len(all_x_chars)\n        frac_num_y = len(re.sub(r\"[^\\d]\", \"\", all_y_chars)) / len(all_y_chars)\n\n        if frac_num_x >= 0.5:\n            x = self.clean(x)\n        else:\n            x = [s.strip() for s in x]\n\n        if frac_num_y >= 0.5:\n            y = self.clean(y)\n        else:\n            y = [s.strip() for s in y]\n\n        return x, y\n\n    def string2preds(self, pred_string: str):\n        \"\"\"\n        Convert the prediction string from Donut to a chart type and x and y values.\n\n        Checks to make sure the special tokens are present and that the x and y values are not empty.\n        Will truncate the list of values to the smaller length of the two lists. This is because the\n        lengths of the x and y values must be the same to earn any points.\n\n        Parameters\n        ----------\n        pred_string : str\n            The prediction string from Donut.\n\n        Returns\n        -------\n        str, list of str, list of str\n            The chart type predicted by Donut, the x values, and the y values.\n        \"\"\"\n        if \"<dot>\" in pred_string:\n            chart_type = \"dot\"\n        elif \"<horizontal_bar>\" in pred_string:\n            chart_type = \"horizontal_bar\"\n        elif \"<vertical_bar>\" in pred_string:\n            chart_type = \"vertical_bar\"\n        elif \"<scatter>\" in pred_string:\n            chart_type = \"scatter\"\n        elif \"<line>\" in pred_string:\n            chart_type = \"line\"\n        else:\n            return \"vertical_bar\", [], []\n\n        if not all([x in pred_string for x in [self.cfg.X_START, \n                                               self.cfg.X_END, \n                                               self.cfg.Y_START, \n                                               self.cfg.Y_END]]):\n            return chart_type, [], []\n\n        pred_string = re.sub(r\"<one>\", \"1\", pred_string)\n\n        x = pred_string.split(self.cfg.X_START)[1].split(self.cfg.X_END)[0].split(\";\")\n        y = pred_string.split(self.cfg.Y_START)[1].split(self.cfg.Y_END)[0].split(\";\")\n\n        if len(x) == 0 or len(y) == 0:\n            return chart_type, [], []\n\n        x, y = self.clean_preds(x, y)\n\n        return chart_type, x, y\n\n    def preprocess(self, examples, processor):\n        \"\"\"\n        Preprocess the images using the defined processor.\n\n        This method converts the input sample (image) into an array and preprocesses it using the processor.\n        If the image is grayscale, it is converted into a 3-channel format.\n\n        Parameters\n        ----------\n        examples : dict\n            Dictionary containing input examples.\n        processor : DonutProcessor\n            The DonutProcessor object for preprocessing.\n\n        Returns\n        -------\n        dict\n            A dictionary containing pixel values of the preprocessed image.\n        \"\"\"\n        pixel_values = []\n\n        for sample in examples[\"image_path\"]:\n            arr = np.array(sample)\n            if len(arr.shape) == 2:\n                arr = np.stack([arr]*3, axis=-1)\n        \n            pixel_values.append(processor(arr, random_padding=True).pixel_values)\n\n        return {\"pixel_values\": torch.tensor(np.vstack(pixel_values))}\n\n    def __load_model(self):\n        \"\"\"Load the pre-trained model and set it to evaluation mode.\"\"\"\n        self.model = VisionEncoderDecoderModel.from_pretrained(self.cfg.model_donut_dir)\n        self.model.eval()\n        self.model.to(self.device)\n        self.decoder_start_token_id = self.model.config.decoder_start_token_id\n        self.processor = DonutProcessor.from_pretrained(self.cfg.model_donut_dir)\n\n    def load_dataset(self):\n        \"\"\"\n        Load the dataset and create a data loader for batching.\n\n        Returns\n        -------\n        list of str\n            List of image IDs.\n        \"\"\"\n        self.__load_model()\n        image_dir = Path(self.cfg.image_path)\n        images = list(image_dir.glob(\"*.jpg\"))\n        self.ds = Dataset.from_dict(\n            {\"image_path\": [str(x) for x in images], \"id\": [x.stem for x in images]}\n        ).cast_column(\"image_path\", ds_img())\n        ids = self.ds[\"id\"]\n        self.ds.set_transform(partial(self.preprocess, processor=self.processor))\n        self.data_loader = DataLoader(\n            self.ds, batch_size=self.cfg.batch_size, shuffle=False\n        )\n        return ids\n\n    def generate_predictions(self):\n        \"\"\"\n        Generate predictions for the input images.\n\n        Returns\n        -------\n        list of str, list of str, list of str\n            List of chart types, x predictions, and y predictions.\n        \"\"\"\n        all_generations = []\n        for batch in tqdm(self.data_loader):\n            pixel_values = batch[\"pixel_values\"].to(self.device)\n            batch_size = pixel_values.shape[0]\n            decoder_input_ids = torch.full(\n                (batch_size, 1),\n                self.decoder_start_token_id,\n                device=pixel_values.device,\n            )\n\n            try:\n                outputs = self.model.generate(\n                    pixel_values,\n                    decoder_input_ids=decoder_input_ids,\n                    max_length=self.cfg.max_length,\n                    early_stopping=True,\n                    pad_token_id=self.processor.tokenizer.pad_token_id,\n                    eos_token_id=self.processor.tokenizer.eos_token_id,\n                    use_cache=True,\n                    num_beams=4,\n                    temperature=.7,\n                    top_k=3,\n                    top_p=.5,\n                    return_dict_in_generate=True,\n                )\n\n                all_generations.extend(self.processor.batch_decode(outputs.sequences))\n\n            except:\n                all_generations.extend([\"\"] * batch_size)\n\n        chart_types, x_preds, y_preds = [], [], []\n        for gen in all_generations:\n            try:\n                chart_type, x, y = self.string2preds(gen)\n                new_chart_type = chart_type\n                x_str = \";\".join(list(map(str, x)))\n                y_str = \";\".join(list(map(str, y)))\n\n            except Exception as e:\n                new_chart_type = self.cfg.PLACEHOLDER_CHART_TYPE\n                x_str = self.cfg.PLACEHOLDER_DATA_SERIES\n                y_str = self.cfg.PLACEHOLDER_DATA_SERIES\n\n            if len(x_str) == 0:\n                x_str = self.cfg.PLACEHOLDER_DATA_SERIES\n            if len(y_str) == 0:\n                y_str = self.cfg.PLACEHOLDER_DATA_SERIES\n\n            chart_types.append(new_chart_type)\n            x_preds.append(x_str)\n            y_preds.append(y_str)\n\n        return chart_types, x_preds, y_preds","metadata":{"execution":{"iopub.status.busy":"2023-05-24T15:08:57.952203Z","iopub.execute_input":"2023-05-24T15:08:57.952631Z","iopub.status.idle":"2023-05-24T15:08:57.991555Z","shell.execute_reply.started":"2023-05-24T15:08:57.952599Z","shell.execute_reply":"2023-05-24T15:08:57.990597Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **3. ResNet50 model**\n\nHere the **ImageCharts** class is defined, which is a subclass of **torch.utils.data.Dataset**. This class is responsible for loading the graphics images and preparing them for the image classification model. It also includes code to load the ResNet50 model and use the model to make predictions on the chart type.","metadata":{}},{"cell_type":"code","source":"class ImageCharts(Dataset):\n    def __init__(self, cfg, df, transforms=None):\n        \"\"\"\n        Initializes the ImageCharts object.\n\n        Parameters\n        ----------\n            cfg (object): Configuration object with necessary information.\n            df (pandas.DataFrame): DataFrame containing the data information.\n            transforms (callable, optional): Transformations to be applied to the images. Default: None.\n        \"\"\"\n        self.cfg = cfg\n        self.df = df\n        self.transforms = transforms\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n    def __len__(self):\n        \"\"\"\n        Returns the size of the dataset.\n\n        Returns:\n            int: Size of the dataset.\n        \"\"\"\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        \"\"\"\n        Returns an item from the dataset at the specified index.\n\n        Args:\n            idx (int): Index of the item.\n\n        Returns:\n            torch.Tensor: Normalized image.\n        \"\"\"\n        name_img = self.df.iloc[idx]['filename']\n        img = cv.imread(name_img)\n        \n        if img is None:\n            raise ValueError(f\"Failed to read image at path: {name_img}\")\n\n        if len(img.shape) == 2:\n            img = cv.cvtColor(img, cv.COLOR_GRAY2RGB)\n        else:\n            img = cv.cvtColor(img, cv.COLOR_BGR2RGB)\n        \n        img = cv.resize(img, (500, 300))\n        img = img.astype(np.float32) / 255.0\n\n        if self.transforms is not None:\n            img = self.transforms(img)\n\n        return img\n\n    def __load_data(self):\n        \"\"\"\n        Loads the data for the dataset.\n        \"\"\"\n        transforms_train = transforms.Compose([\n            transforms.ToTensor(),\n            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n        ])\n        self.transforms = transforms_train\n\n    def predict(self):\n        \"\"\"\n        Performs predictions based on the dataset data.\n\n        Returns:\n            list: List containing the predictions for each image.\n        \"\"\"\n        self.__load_data()\n        test_data_loader = DataLoader(self, batch_size=self.cfg.batch_size, shuffle=True)\n\n        model = torchvision.models.resnet50(pretrained=False)\n        model.fc = nn.Linear(in_features=2048, out_features=self.cfg.num_classes)\n\n        model.load_state_dict(torch.load(self.cfg.model_ResNet50_dir))\n        model.eval()\n\n        all_outputs = []\n        with torch.no_grad():\n            for inputs in tqdm(test_data_loader):\n                inputs = inputs.to(self.device)\n                model = model.to(self.device)\n                outputs = model(inputs)\n                _, preds = torch.max(outputs, 1)\n                all_outputs.extend(preds.cpu().numpy())\n\n        return all_outputs","metadata":{"execution":{"iopub.status.busy":"2023-05-24T15:08:57.995041Z","iopub.execute_input":"2023-05-24T15:08:57.995666Z","iopub.status.idle":"2023-05-24T15:08:58.009861Z","shell.execute_reply.started":"2023-05-24T15:08:57.995634Z","shell.execute_reply":"2023-05-24T15:08:58.008699Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **4. Predict with DONUT model**\nThe DonutDataProcessor is invoked to load the dataset and generate predictions for the data series in the charts.","metadata":{}},{"cell_type":"code","source":"processor = DonutDataProcessor(CFG)\nids_data = processor.load_dataset()\n\nchart_types, x_preds, y_preds = processor.generate_predictions()","metadata":{"execution":{"iopub.status.busy":"2023-05-24T15:08:58.011321Z","iopub.execute_input":"2023-05-24T15:08:58.011665Z","iopub.status.idle":"2023-05-24T15:09:27.673340Z","shell.execute_reply.started":"2023-05-24T15:08:58.011633Z","shell.execute_reply":"2023-05-24T15:09:27.672134Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **5. Prediction with ResNet50 model**\n\nThe ImageCharts is instantiated to load the image dataset and make predictions about the chart type.","metadata":{}},{"cell_type":"code","source":"df_filename = pd.DataFrame({\"filename\": [os.path.join(CFG.image_path, f\"{id_}.jpg\") for id_ in ids_data]})\n\ndataset = ImageCharts(cfg=CFG, df=df_filename)\n\npredictions = dataset.predict()","metadata":{"execution":{"iopub.status.busy":"2023-05-24T15:09:27.674729Z","iopub.execute_input":"2023-05-24T15:09:27.675177Z","iopub.status.idle":"2023-05-24T15:09:29.437288Z","shell.execute_reply.started":"2023-05-24T15:09:27.675144Z","shell.execute_reply":"2023-05-24T15:09:29.436295Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **6. Combination of models**\n\nThe predictions of the two models are combined. If both models agree on the chart type, that type is selected. **If they disagree, one of the two types is chosen at random**.","metadata":{}},{"cell_type":"code","source":"def select_charts(list1, list2, mapping_dict):\n    \"\"\"\n    Selects charts from list1 and list2 based on their values, with the option to choose randomly if the values differ.\n    \n    Parameters\n    ----------\n        list1 (list): The first list of values.\n        list2 (list): The second list of values.\n        mapping_dict (dict): A dictionary that maps values in list2 to corresponding values in list1.\n    \n    Returns:\n        list: The selected charts based on the comparison of values in list1 and list2.\n    \"\"\"\n    transformed_list2 = [mapping_dict[value] for value in list2]\n\n    selected_charts = []\n\n    for value1, value2 in zip(list1, transformed_list2):\n        if value1 == value2:\n            selected_charts.append(value1)\n        else:\n            selected_charts.append(random.choice([value1, value2]))\n\n    return selected_charts","metadata":{"execution":{"iopub.status.busy":"2023-05-24T15:09:29.438796Z","iopub.execute_input":"2023-05-24T15:09:29.439780Z","iopub.status.idle":"2023-05-24T15:09:29.446637Z","shell.execute_reply.started":"2023-05-24T15:09:29.439746Z","shell.execute_reply":"2023-05-24T15:09:29.445528Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"new_chart_types = select_charts(chart_types, predictions, CFG.label_map)","metadata":{"execution":{"iopub.status.busy":"2023-05-24T15:09:29.447998Z","iopub.execute_input":"2023-05-24T15:09:29.448416Z","iopub.status.idle":"2023-05-24T15:09:29.456764Z","shell.execute_reply.started":"2023-05-24T15:09:29.448385Z","shell.execute_reply":"2023-05-24T15:09:29.455632Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub_df = pd.DataFrame(\n    data={\n        \"id\": [f\"{id_}_x\" for id_ in ids_data] + [f\"{id_}_y\" for id_ in ids_data],\n        \"data_series\": x_preds + y_preds,\n        \"chart_type\": new_chart_types * 2,\n    }\n)","metadata":{"execution":{"iopub.status.busy":"2023-05-24T15:09:29.477594Z","iopub.execute_input":"2023-05-24T15:09:29.477934Z","iopub.status.idle":"2023-05-24T15:09:29.484871Z","shell.execute_reply.started":"2023-05-24T15:09:29.477905Z","shell.execute_reply":"2023-05-24T15:09:29.484151Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub_df.to_csv('submission.csv', index = False)","metadata":{"execution":{"iopub.status.busy":"2023-05-24T15:09:29.486090Z","iopub.execute_input":"2023-05-24T15:09:29.487100Z","iopub.status.idle":"2023-05-24T15:09:29.496455Z","shell.execute_reply.started":"2023-05-24T15:09:29.487069Z","shell.execute_reply":"2023-05-24T15:09:29.495857Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **7. Conclusion**\n\nThe use of a **hybrid model** provided the necessary **flexibility** to take advantage of the distinct strengths of two different models. The Donut model was used to generate initial predictions of the graph types and data series, while an image classification model (ResNet50) was employed to predict the graph type directly from the image.","metadata":{}}]}