theovercomer8 / captionr

GIT/BLIP/CLIP Caption tool
MIT License
138 stars 15 forks source link

edit captions: added utf8 argument to save, added caption saved messa… #19

Open younyokel opened 1 year ago

younyokel commented 1 year ago

…ge and fixed img quality typo

younyokel commented 1 year ago

The colab file underwent some uneeeded changes so I'll just share it here

import ipywidgets as widgets
from io import BytesIO
#@markdown #Edit Captions

#@markdown - Open a tool to view/edit captions.

paths=""
out=""
widgets_l=""
def Caption(path):
    name = os.path.splitext(os.path.basename(path))[0]
    ext=os.path.splitext(os.path.basename(path))[-1][1:]
    if ext=="jpg" or "JPG":
      ext="JPEG"      

    if os.path.exists(folder_path+"/"+name + '.' + extension):
      with open(folder_path+"/"+name + '.' + extension, 'r', encoding='utf-8') as f:
          text = f.read()
    else:
      with open(folder_path+"/"+name + '.' + extension, 'w', encoding='utf-8') as f:
          f.write("")
          with open(folder_path+"/"+name + '.' + extension, 'r', encoding='utf-8') as f:
              text = f.read()   

    img=Image.open(os.path.join(folder_path,path))
    aspect_ratio = img.width / img.height
    max_size = 460  # Set the maximum size here

    # Calculate new dimensions while maintaining aspect ratio
    width, height = img.size
    aspect_ratio = width / height

    if width > height:
        new_width = max_size
        new_height = int(new_width / aspect_ratio)
    else:
        new_height = max_size
        new_width = int(new_height * aspect_ratio)

    img = img.resize((new_width, new_height))
    image_bytes = BytesIO()
    img.save(image_bytes, format=ext, quality=80)
    image_bytes.seek(0)
    image_data = image_bytes.read()

    image = widgets.Image(
        value=image_data,
        width=new_width,
        height=new_height
    )
    text_area = widgets.Textarea(value=text, description='', disabled=False, layout={'width': '300px', 'height': '120px'})

    # create a label to display the save status
    save_status = widgets.Label(value='')

    def update_text(text):
        save_status.value = ''
        with open(folder_path+"/"+name + '.' + extension, 'w') as f:
            f.write(text)
        # update the save status label
        save_status.value = 'Image saved!'

    button = widgets.Button(description='Save', button_style='success')
    button.on_click(lambda b: update_text(text_area.value))
    text_area.observe(lambda change: setattr(save_status, 'value', ''), names='value')

    # change the layout of the widgets
    return widgets.VBox([
        widgets.HBox([image, widgets.VBox([text_area, button, save_status])])
    ])

paths = os.listdir(folder_path)
paths = list(filter(lambda fn: not fn.endswith(extension), paths))

widgets_l = widgets.Select(options=["Select an image to caption"]+paths, rows=25)

out = widgets.Output()

def click(change):
    with out:
        out.clear_output()
        display(Caption(change.new))

widgets_l.observe(click, names='value')
display(widgets.HBox([widgets_l, out]))