really dumb ui for looking at grounding files from kosmos2

This commit is contained in:
Victor Hall 2024-03-03 15:13:08 -05:00
parent 8bab958c68
commit a55ab816af
1 changed files with 98 additions and 0 deletions

98
scripts/grounding_ui.py Normal file
View File

@ -0,0 +1,98 @@
import tkinter as tk
from tkinter import filedialog, Listbox
from PIL import Image, ImageTk
import cv2
import numpy as np
import json
import os
import ast
class ImageBoundingBoxApp:
def __init__(self, root, folder_path):
self.root = root
self.folder_path = folder_path
self.images = self.load_images(folder_path)
self.setup_ui()
def load_images(self, folder_path):
images = [f for f in os.listdir(folder_path) if f.endswith(('.jpg', '.png', '.webp'))]
return images
def setup_ui(self):
self.listbox = Listbox(self.root)
self.listbox.pack(side="left", fill="y")
self.canvas = tk.Canvas(self.root, bg='white')
self.canvas.pack(side="right", expand=True, fill="both")
for image in self.images:
self.listbox.insert(tk.END, image)
self.listbox.bind("<<ListboxSelect>>", self.on_select)
def on_select(self, event):
w = event.widget
index = int(w.curselection()[0])
image_name = w.get(index)
self.display_image_with_bboxes(image_name)
def display_image_with_bboxes(self, image_name):
img_path = os.path.join(self.folder_path, image_name)
ent_path = img_path.rsplit('.', 1)[0] + '.ent'
# Load image
image = cv2.imread(img_path)
image_h, image_w, _ = image.shape
font_scale = 0.5 * (image_h + image_w) / 1000
thickness = max(int(0.5 * (image_h + image_w) / 500),1)
print(f"Font Scale: {font_scale}, Thickness: {thickness}")
# Read .ent file
if os.path.exists(ent_path):
with open(ent_path, 'r') as f:
entities = ast.literal_eval(f.read())
for entity_name, _, bboxes in entities:
for (x1_norm, y1_norm, x2_norm, y2_norm) in bboxes:
orig_x1, orig_y1, orig_x2, orig_y2 = int(x1_norm * image_w), int(y1_norm * image_h), int(x2_norm * image_w), int(y2_norm * image_h)
color = (0, 255, 0) # Green color for bounding box
cv2.rectangle(image, (orig_x1, orig_y1), (orig_x2, orig_y2), color, 2)
cv2.putText(image, entity_name, (orig_x1, orig_y1 + int(30*font_scale)), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, thickness)
# Convert to PIL Image to display in Tkinter
image_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
# Resize image to fit the canvas size while maintaining aspect ratio
canvas_width = self.canvas.winfo_width()
canvas_height = self.canvas.winfo_height()
img_width, img_height = image_pil.size
# Calculate the proper scaling factor
scale_w = canvas_width / img_width
scale_h = canvas_height / img_height
scale = min(scale_w, scale_h)
# Calculate new image size
new_width = int(img_width * scale)
new_height = int(img_height * scale)
# Resize the image with new dimensions
image_resized = image_pil.resize((new_width, new_height), Image.Resampling.LANCZOS)
image_tk = ImageTk.PhotoImage(image_resized)
# Clear the canvas and display the resized image
self.canvas.delete("all") # Clear previous image/text
self.canvas.create_image((canvas_width - new_width) // 2, (canvas_height - new_height) // 2, anchor="nw", image=image_tk)
self.canvas.image = image_tk # Keep a reference!
def main():
root = tk.Tk()
root.title("Image Bounding Box Viewer")
folder_path = filedialog.askdirectory(title="Select Folder with Images")
app = ImageBoundingBoxApp(root, folder_path)
root.mainloop()
if __name__ == "__main__":
main()