helper_functions.py

from io import BytesIO
from PIL import Image
from matplotlib import pyplot as plt
import numpy as np
import requests
import torch
import os
import glob
import pandas as pd
import xml.etree.ElementTree as ET

def save_images(folder, search_term, count=10):
    if not os.path.exists(folder):
        os.mkdir(folder)

#    SEARCH_URL = "https://huggingface.co/api/experimental/images/search"
#    params = {"q": search_term, "license": "public", "imageType": "photo", "count": count}
    SEARCH_URL = "https://pixabay.com/api/"
    params = {"q": search_term.replace(' ', '+'), 
              "image_type": "photo", 
              "per_page": count, 
              "page": 1, 
              "key": os.environ["PIXABAY_KEY"]}

    resp = requests.get(SEARCH_URL, params=params)
    if resp.status_code == 200:
#        content = resp.json()['value']
#        urls = [img['thumbnailUrl'] for img in content]
        content = resp.json()['hits']
        urls = [img['previewURL'] for img in content]

        folder = os.path.join(folder, search_term)
        if not os.path.exists(folder):
            os.mkdir(folder)

        i = 0
        for url in urls:
            try:
                img = get_image_from_url(url)
                fname = os.path.join(folder, f'{i}.jpg')
                img.save(fname)
                i += 1
            except Exception:
                pass
        print(f'Retrieved {i} images for {search_term}')
    else:
        print(f'Failed to retrieve URLs for {search_term}')

def show(imgs, titles=None):
    import torchvision.transforms.functional as tF
    
    if not isinstance(imgs, list):
        imgs = [imgs]
    fig, axs = plt.subplots(ncols=len(imgs), squeeze=False)
    for i, img in enumerate(imgs):
        if isinstance(img, torch.Tensor):
            img = img.detach()
            img = tF.to_pil_image(img)
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
        if titles is not None:
            if i < len(titles):
                axs[0, i].set_title(titles[i])
    return fig

def get_image_from_url(url, headers=None):
    resp = requests.get(url, headers=headers)
    resp.raise_for_status()
    img = Image.open(BytesIO(resp.content))
    return img

def xml_to_csv(path):
    """Iterates through all .xml files (generated by labelImg) in a given directory and combines
    them in a single Pandas dataframe.

    Parameters:
    ----------
    path : str
        The path containing the .xml files
    Returns
    -------
    Pandas DataFrame
        The produced dataframe
    """

    xml_list = []
    for xml_file in glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        filename = root.find('filename').text
        width = int(root.find('size').find('width').text)
        height = int(root.find('size').find('height').text)
        for member in root.findall('object'):
            bndbox = member.find('bndbox')
            value = (filename,
                     width,
                     height,
                     member.find('name').text,
                     int(bndbox.find('xmin').text),
                     int(bndbox.find('ymin').text),
                     int(bndbox.find('xmax').text),
                     int(bndbox.find('ymax').text),
                     )
            xml_list.append(value)
    column_name = ['filename', 'width', 'height',
                   'class', 'xmin', 'ymin', 'xmax', 'ymax']
    xml_df = pd.DataFrame(xml_list, columns=column_name)
    return xml_df

def get_summarizer(model, tokenizer):
	def summarizer(
	    texts,
	    max_length=128,
	    min_length=30,
	    num_beams=4,
	    do_sample=False,
	    truncation=True,
	):
	    single_input = isinstance(texts, str)
	    if single_input:
    		texts = [texts]

	    batch = tokenizer(
    		texts,
    		return_tensors="pt",
    		padding=True,
    		truncation=truncation,
	    ).to(model.device)

	    model.eval()
	    with torch.inference_mode():
		output_ids = model.generate(
		    **batch,
		    max_new_tokens=max_length,
		    min_new_tokens=min_length,
		    num_beams=num_beams,
		    do_sample=do_sample,
		)

	    summaries = tokenizer.batch_decode(
    		output_ids,
    		skip_special_tokens=True,
    		clean_up_tokenization_spaces=True,
	    )
	    result = [{"summary_text": s} for s in summaries]
	    return result[0] if single_input else result
	
    return summarizer