-
-
Notifications
You must be signed in to change notification settings - Fork 116
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
really dumb ui for looking at grounding files from kosmos2
- Loading branch information
1 parent
8bab958
commit a55ab81
Showing
1 changed file
with
98 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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() |