utils.py

# import statements
import time
import io
import os
import csv
import hashlib
import re
from tqdm import tqdm
import requests
from urllib.request import urlopen
from urllib.parse import urljoin, urlparse
from bs4 import BeautifulSoup
from joblib import Parallel, delayed
from tqdm import tqdm

import logging

def filter_text(tag_text):
    """
    This returns True if the number of words in the text is less than 3
    """
    if len(tag_text.split(' ')) > 2:
        return False
    else:
        return True

def clean_text(extracted_text):
    """
    This function cleans th etext for any extra spaces or newlines.
    """
    return ' '.join(extracted_text.split())

def write_list(list_to_write, fname):
    import pickle
    """
    Write python list to a pickle file.
    """
    # Save all the website urls extracted
    with open(fname, 'wb') as out_file:
        pickle.dump(list_to_write, out_file)

def write_dict(folder_path, dict_to_write, fname, fieldnames):
    """
    Write a dictionary to a file.
    """
    output_fname = os.path.join(folder_path, fname)
    print(f'Writing : {output_fname}')
    with open(output_fname, 'wt') as out_file:
        writer = csv.DictWriter(out_file, fieldnames=fieldnames, delimiter='\t')
        for key, value in dict_to_write.items():
            data = {fieldnames[0]: key, fieldnames[1]: value}
            writer.writerow(data)

def collect_web_urls(file_path):
    """
    Reads a url text file for urls and extracts all the web urls saved.
    """
    logging.info("Reading the url text file.")
    urls = set()
    with open(file_path, 'r') as urls_file:
        url_list = urls_file.readlines()
        for url in url_list:
            urls.add(url.strip())
        
    return list(urls)


def process_web_url(web_url):
    """
    Function to process each web URL and extract relevant image data.
    """
    webpage_image_urls, webpage_image_alt_text, webpage_text_above, webpage_text_below, webpage_image_class_names = fetch_images(web_url)
    time.sleep(1)  # To avoid overwhelming the server

    # Store the results in a dictionary
    result = {
        'webpage_image_urls': webpage_image_urls,
        'webpage_image_alt_text': webpage_image_alt_text,
        'webpage_text_above': webpage_text_above,
        'webpage_text_below': webpage_text_below,
        'webpage_image_class_names': webpage_image_class_names
    }

    return result

def extract_images(website_urls, folder_path):
    """
    This method goes over each url and extracts all the images and associated text and
    creates 5 files:
    1. image_urls.txt (contains image source urls for all the images extracted)
    2. img_url_to_captions_new.csv (image urls and corresponding alt text)
    3. img_url_to_text_above_new.csv (image urls and corresponding preceding text)
    4. img_url_to_text_below_new.csv (image urls and corresponding succeeding text)
    5. image_url_to_image_class_names_new.csv (image-urls and the corresponding img tag class name)
    """
    if not os.path.exists(folder_path):
        os.mkdir(folder_path)

    logging.info("Processing all web_urls.")

    # Use Parallel and delayed to parallelize the fetching process
    results = Parallel(n_jobs=-1)(delayed(process_web_url)(web_url) for web_url in tqdm(website_urls))

    # Initialize dictionaries and lists
    image_url_dict = {}
    image_urls = []
    image_alt_text = []
    text_above = []
    text_below = []
    image_class_names = []
    img_url_to_caption = {}
    img_url_to_text_above = {}
    img_url_to_text_below = {}
    img_url_to_class_name = {}

    # Process each result
    for result in results:
        webpage_image_urls = result['webpage_image_urls']
        webpage_image_alt_text = result['webpage_image_alt_text']
        webpage_text_above = result['webpage_text_above']
        webpage_text_below = result['webpage_text_below']
        webpage_image_class_names = result['webpage_image_class_names']

        # Remove any duplicate images
        idx = 0
        for url in webpage_image_urls:
            if url in image_url_dict:
                idx += 1
                continue
            else:
                image_url_dict[url] = 1
                image_urls.append(url)
                image_alt_text.append(webpage_image_alt_text[idx])
                text_above.append(webpage_text_above[idx])
                text_below.append(webpage_text_below[idx])
                image_class_names.append(webpage_image_class_names[idx])

                # Add to img_url dictionaries
                img_url_to_caption[url] = webpage_image_alt_text[idx]
                img_url_to_class_name[url] = webpage_image_class_names[idx]
                img_url_to_text_above[url] = webpage_text_above[idx]
                img_url_to_text_below[url] = webpage_text_below[idx]
                idx += 1

    # Dump all image_urls
    write_list(image_urls, os.path.join(folder_path, 'ui_images.p'))

    # Write the image url dictionaries
    write_dict(folder_path, img_url_to_caption, 'ui_alt_texts.csv', ['Image_Url', 'Image_Alt_Text'])
    write_dict(folder_path, img_url_to_text_above, 'ui_instructions_preceding.csv', ['Image_Url', 'Text_Above'])
    write_dict(folder_path, img_url_to_text_below, 'ui_instructions_succeeding.csv', ['Image_Url', 'Text_Below'])
    # Classnames can be useful in filtering noisy images like ads etc.
    write_dict(folder_path, img_url_to_class_name, 'ui_image_url_to_image_class_names.csv', ['Image_Url', 'Class_Name'])



def fetch_images(website_url):
    """
    This method extracts the data on a web url and then extracts all images present on the
    webpage. For each img tag, it then extracts the class name for it and the preceding
    and succeeding text inside a ul, l, p or div tag. We also apply the length filter on the
    text extracted, i.e. if the number of words is less than 3, we extract more preceding text.
    """
    try_count = 0
    while True:
        try:
            response = requests.get(website_url, timeout=10)
            soup = BeautifulSoup(response.text, 'html.parser')
            break
        except Exception:
            try_count += 1
        if try_count > 5:
            with open('error_log.txt', 'a') as error_log:
                error_log.write(f"ERROR - Could not retrieve {website_url}\n")
            return [],[],[],[],[]    
        
    img_tags = soup.find_all('img', src=True, alt=True)
    
    urls = [img['src'] for img in img_tags]
    
    # Extract all image captions. If img does not have alt text, append empty string.
    image_alt_text = [img['alt'] for img in img_tags]
    
    # Format urls to get list of all image urls from the webpage
    image_urls = []  
    for url in urls:
        if 'http' not in url:
            url = '{}{}'.format(website_url, url)
        image_urls.append(url)

    # Extract text
    text_above = []
    text_below = []
    img_class_names = []
    for img_tag in img_tags:
        if img_tag.has_attr('class'):
            img_class_names.append(img_tag['class'])
        else:
            img_class_names.append("")
        prev_not_found = True
        current_tag = img_tag
        while prev_not_found:
            prev_tag = current_tag.previous_element
            if not prev_tag:
                break
            # EXTRACT TEXT FROM WHATEVER THE PARENT TAG IS (<p>, <div>, <ul>, <l> etc.)
            if prev_tag.name in ['p', 'div', 'ul', 'l']:
                if not filter_text(prev_tag.getText()): #TODO: make sure total number of words is > 2 len(prev_tag.getText()) >= 5 and (not filter_p_tag(prev_tag.getText()))
                    prev_not_found = False
                else:
                    current_tag = prev_tag  
            else:
                current_tag = prev_tag

        next_not_found = True
        current_tag = img_tag
        while next_not_found:
            next_tag = current_tag.next_element
            if not next_tag:
                break
            # EXTRACT TEXT FROM WHATEVER THE PARENT TAG IS (<p>, <div>, <ul>, <l> etc.)
            if next_tag.name in ['p', 'div', 'ul', 'l']:
                if not filter_text(next_tag.getText()): #TODO: make sure total number of words is > 2
                    next_not_found = False
                else:
                    current_tag = next_tag
            else:
                current_tag = next_tag

        if prev_not_found:
            text_above.append('')
        else:
            text_above.append(clean_text(prev_tag.getText()))
        
        if next_not_found:
            text_below.append('')
        else:

            text_below.append(clean_text(next_tag.getText()))


    return image_urls, image_alt_text, text_above, text_below, img_class_names


def process_images_in_parallel(image_urls, query_folder_path, img_url_to_img_id, writer):
    """
    This method processes image URLs in parallel and writes the results using the provided writer.
    """
    def process_image_url(image_url, query_folder_path, img_url_to_img_id):
        """
        Function to process a single image URL.
        """
        img_fname = ''
        if image_url not in img_url_to_img_id:
            img_fname = persist_image(query_folder_path, image_url)
            if img_fname == '':
                return None  # Skip if image could not be persisted
            img_url_to_img_id[image_url] = img_fname
        return {"Image_Url": image_url, "Image_Name": img_fname}
    # Use Parallel and delayed to parallelize the processing of image URLs
    results = Parallel(n_jobs=-1)(delayed(process_image_url)(image_url, query_folder_path, img_url_to_img_id) for image_url in tqdm(image_urls))

    # Write the results to the CSV using the writer
    for data in results:
        if data is not None:  # Skip None results
            writer.writerow(data)

def download_images(image_urls, query_folder_path):
    """
    Traverses through the list of validated image urls and downloads each of them 
    and saves the image-url to image path map to image_urls_processed.csv file
    """
    done_image_urls_fname = os.path.join(query_folder_path, 'image_urls_processed.csv')

    img_url_to_img_id = {}
    if os.path.exists(done_image_urls_fname):
        with open(done_image_urls_fname,newline='') as processed_file:
            file_reader = csv.reader(processed_file, delimiter='\t')
            for row in file_reader:
                img_url = urlparse(row[0].strip()).geturl()
                img_url_to_img_id[img_url] = row[1]

    print('Downloading all images now...')
    with open(done_image_urls_fname, 'w') as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=["Image_Url", "Image_Name"], delimiter='\t')
        process_images_in_parallel(image_urls, query_folder_path, img_url_to_img_id, writer)

def persist_image(folder_path:str, url:str):
    """
    Downloads an image using the image src url and returns the image file-path.
    If the image cannot be downloaded, it returns an empty string.
    """
    try :
        image_content = requests.get(url, timeout=10).content
    except Exception as e:
        #print(f"ERROR - Could not download {url} - {e}")
        return ''

    fname = os.path.join(folder_path,hashlib.sha1(image_content).hexdigest()[:10] + '.jpg')
    try:
        img_file = open(fname, "wb")
        img_file.write(image_content)
        img_file.close()
    except Exception as e:
        print(f"ERROR - Could not save {url} - {e}")
        return ''
    
    return fname