counterfactual-world-models / gradio_app_intervention.py
rahulvenkk
app.py updated
6dfcb0f
raw
history blame
10.1 kB
import cv2
import numpy as np
import gradio as gr
import cwm.utils as utils
# Points color and arrow properties
arrow_color = (0, 255, 0) # Green color for all arrows
dot_color = (0, 255, 0) # Green color for the dots at start and end
dot_color_fixed = (255, 0, 0) # Red color for zero-length vectors
thickness = 4 # Thickness of the arrow
tip_length = 0.3 # The length of the arrow tip relative to the arrow length
dot_radius = 10 # Radius for the dots
dot_thickness = -1 # Thickness for solid circle (-1 fills the circle)
from PIL import Image
import torch
#load model
from cwm.model.model_factory import model_factory
from timm.data.constants import (IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Load CWM 3-frame model (automatically download pre-trained checkpoint)
model = model_factory.load_model('vitb_8x8patch_3frames').to(device)
model.requires_grad_(False)
model.eval()
model = model.to(torch.float16)
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch
from PIL import Image
import numpy as np
def draw_arrows_matplotlib(img, selected_points, zero_length):
"""
Draw arrows on the image using matplotlib for better quality arrows and dots.
"""
fig, ax = plt.subplots()
ax.imshow(img)
for i in range(0, len(selected_points), 2):
start_point = selected_points[i]
end_point = selected_points[i + 1]
if start_point == end_point or zero_length:
# Draw a dot for zero-length vectors or if only one point is clicked
ax.scatter(start_point[0], start_point[1], color='red', s=100) # Red dot for zero-length vector
else:
# Draw arrows
arrow = FancyArrowPatch((start_point[0], start_point[1]), (end_point[0], end_point[1]),
color='green', linewidth=2, arrowstyle='->', mutation_scale=15)
ax.add_patch(arrow)
# Optionally, draw a small circle (dot) at the start and end points
ax.scatter(start_point[0], start_point[1], color='green', s=100) # Green dot at start
ax.scatter(end_point[0], end_point[1], color='green', s=100) # Green dot at end
# Save the image to a numpy array
fig.canvas.draw()
img_array = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
img_array = img_array.reshape(fig.canvas.get_width_height()[::-1] + (3,))
plt.close(fig)
return img_array
with gr.Blocks() as demo:
with gr.Row():
gr.Markdown('''# Generate interventions!πŸš€
Upload an image and click to select the start and end points for arrows. Dots will be shown at the beginning and end of each arrow. You can also create zero-length vectors (just a dot) by enabling the toggle below.
''')
# Annotating arrows on an image
with gr.Tab(label='Image'):
with gr.Row():
with gr.Column():
# Input image
original_image = gr.State(value=None) # store original image without arrows
input_image = gr.Image(type="numpy", label="Upload Image")
# Annotate arrows
selected_points = gr.State([]) # store points
zero_length_toggle = gr.Checkbox(label="Enable zero-length vectors", value=False) # Toggle for zero-length vectors
with gr.Row():
gr.Markdown('Click on the image to select the start and end points for each arrow. If zero-length vectors are enabled, clicking once will draw a dot.')
undo_button = gr.Button('Undo last action')
clear_button = gr.Button('Clear All')
# Run model button
run_model_button = gr.Button('Run Model')
# Show the image with the annotated arrows
with gr.Tab(label='Intervention'):
output_image = gr.Image(type='numpy')
def resize_to_square(img, size=512):
img = Image.fromarray(img)
img = img.resize((size, size))
return np.array(img)
# Store the original image and resize to square size once uploaded
def store_img(img):
resized_img = resize_to_square(img) # Resize the uploaded image to a square
print(f"Image uploaded with shape: {resized_img.shape}")
return resized_img, resized_img, []
input_image.upload(store_img, [input_image], [input_image, original_image, selected_points])
# Get points and draw arrows or zero-length vectors based on the toggle
def get_point(img, sel_pix, zero_length, evt: gr.SelectData):
sel_pix.append(evt.index) # Append the point's location (coordinates)
# Zero-length vector case: Draw a single dot at the clicked point
if zero_length:
point = sel_pix[-1] # Last point clicked
cv2.circle(img, point, dot_radius, dot_color_fixed, dot_thickness) # Draw a dot at the point
sel_pix.append(evt.index)
else:
# Regular case: two clicks for an arrow
# Check if this is the first point (start point for the arrow)
if len(sel_pix) % 2 == 1:
# Draw a dot at the start point to give feedback
start_point = sel_pix[-1] # Last point is the start
cv2.circle(img, start_point, dot_radius, dot_color, dot_thickness)
# Check if two points have been selected (start and end points for an arrow)
if len(sel_pix) % 2 == 0:
# Draw an arrow between the last two points
start_point = sel_pix[-2] # Second last point is the start
end_point = sel_pix[-1] # Last point is the end
# Draw arrow
cv2.arrowedLine(img, start_point, end_point, arrow_color, thickness, tipLength=tip_length)
# Draw a dot at the end point
cv2.circle(img, end_point, dot_radius, dot_color, dot_thickness)
return img if isinstance(img, np.ndarray) else np.array(img)
input_image.select(get_point, [input_image, selected_points, zero_length_toggle], [input_image])
# Undo the last selected action
def undo_arrows(orig_img, sel_pix, zero_length):
temp = orig_img.copy()
# if zero_length:
# # Undo the last zero-length vector (just the last dot)
# if len(sel_pix) >= 1:
# sel_pix.pop() # Remove the last point
# else:
if len(sel_pix) >= 2:
sel_pix.pop() # Remove the last end point
sel_pix.pop() # Remove the last start point
# Redraw all remaining arrows and dots
for i in range(0, len(sel_pix), 2):
start_point = sel_pix[i]
end_point = sel_pix[i + 1]
if start_point == end_point:
# Zero-length vector: Draw a dot
color = dot_color_fixed
else:
cv2.arrowedLine(temp, start_point, end_point, arrow_color, thickness, tipLength=tip_length)
color = arrow_color
# Draw arrow
# Draw dots at start and end points
cv2.circle(temp, start_point, dot_radius, color, dot_thickness)
cv2.circle(temp, end_point, dot_radius, color, dot_thickness)
# If there is an odd number of points (e.g., only a start point), draw a dot for it
if len(sel_pix) == 1:
start_point = sel_pix[0]
cv2.circle(temp, start_point, dot_radius, dot_color, dot_thickness)
return temp if isinstance(temp, np.ndarray) else np.array(temp)
undo_button.click(undo_arrows, [original_image, selected_points, zero_length_toggle], [input_image])
# Clear all points and reset the image
def clear_all_points(orig_img, sel_pix):
sel_pix.clear() # Clear all points
return orig_img # Reset image to original
clear_button.click(clear_all_points, [original_image, selected_points], [input_image])
# Dummy model function to simulate running a model
def run_model_on_points(points, input_image, original_image):
H = input_image.shape[0]
W = input_image.shape[1]
factor = 224/H
# Example: pretend the model processes points and returns a simple transformation on the image
points = torch.from_numpy(np.array(points).reshape(-1, 4)) * factor
points = points[:, [1, 0, 3, 2]]
print(points)
img = Image.fromarray(original_image)
img = img.resize((224, 224))
img = np.array(img)
np.save("img.npy", original_image)
img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
img = img[None]
# reshape image to [B, C, T, H, W], C = 3, T = 3 (3-frame model), H = W = 224
x = img[:, :, None].expand(-1, -1, 3, -1, -1).to(torch.float16)
# Imagenet-normalize the inputs (standardization)
x = utils.imagenet_normalize(x).to(device)
with torch.no_grad():
counterfactual = model.get_counterfactual(x, points)
counterfactual = counterfactual.squeeze()
counterfactual = counterfactual.clamp(0, 1).permute(1,2,0).detach().cpu().numpy()
# for i in range(0, len(points), 2):
# # Draw rectangles on the points as model output example
# cv2.rectangle(processed_image, points[i], points[i + 1], (255, 0, 0), 3)
return counterfactual
# Run model when the button is clicked
run_model_button.click(run_model_on_points, [selected_points, input_image, original_image], [output_image])
# Launch the app
demo.queue().launch(inbrowser=True, share=True)