rename to fit

app.py

@@ -0,0 +1,308 @@
+import sys
+sys.path.insert(0, "./SMT")
+
+from smt_trainer import SMT_Trainer
+from smt_model.modeling_smt import SMTModelForCausalLM
+
+import torch
+import gradio as gr
+import numpy as np
+import pandas as pd
+import cv2
+
+from math import sqrt
+
+CA_layers = list()
+
+colors = [  (128,   0,   0),
+			(128,  64,   0),
+			(128, 128,   0),
+			(  0, 128,   0),
+			(  0, 128, 128),
+			(  0,  64, 128),
+			(  0,   0, 128),
+			(128,   0, 128),
+			(128,   0,   0)
+			]
+
+def contrast(elem):
+	return elem!=0
+
+def overlay(background:np.ndarray, overlay:np.ndarray, alpha=1):
+	"""
+	:param background: BGR image (np.uint8)
+	:param overlay: BGRA image (np.uint8)
+	:param alpha: Transparency of overlay over background
+	
+	returns BGR image of combined images (np.float32)
+	"""
+
+	# add alpha channel to background
+	background = np.concatenate([background, np.full([*background.shape[:2], 1], 1.0)], axis=-1 )
+
+	# normalize overlay alpha channel from 0-255 to 0.-1.
+	alpha_background = 1.0
+	alpha_overlay = overlay[:,:,3] / 255.0 * alpha
+
+	for channel in range(3):
+
+		background[:,:,channel] = 	alpha_overlay * overlay[:,:,channel] + \
+									alpha_background * background[:,:,channel] * ( 1 - alpha_overlay )
+
+	background[:,:,3] = ( 1 - ( 1 - alpha_overlay ) * ( 1 - alpha_background ) ) * 255
+
+	# ignore alpha channel because gradio doesnt care
+	# also divide by 255 because somehow it needs a float image even though it gives int images
+	return (background[:,:,:3]/255.0).astype(np.float32)
+
+def generate_CA_images(token_idx, image, multiplier=1):
+
+	global CA_layers
+
+	CA_final_images = []
+
+	# resize to fit input image (value in 0-1)
+	masks = [ cv2.resize(CA_layers[layer_idx][token_idx],
+							interpolation=cv2.INTER_NEAREST,
+							dsize=(image.shape[1], image.shape[0])) for layer_idx in range(0, len(CA_layers)) ]
+
+	for i,mask in enumerate(masks):
+
+		# apply multiplier
+		mask *= multiplier
+		
+		# normalize values above 1
+		max_pixel = np.max(mask)
+		if max_pixel > 1:
+			mask /= max_pixel
+
+		# (convert to values in 0-255)
+		mask = np.round(mask*255.0).astype(np.uint8)
+
+		# add singleton dimension as channel
+		mask = np.expand_dims(mask, axis=-1)
+
+		# base color + transparency mask = BGRA
+		ca = np.concatenate( (np.full(shape=image.shape, fill_value=colors[i]), mask ), axis=-1)
+
+		CA_final_images.append(overlay(image, ca))
+
+	return CA_final_images
+
+def make_predictions(checkpoint, input_image, input_type:int):
+
+	global CA_layers
+
+	# take from huggingface
+	if input_type == 0:
+		# TODO this doesnt work because the HuggingFace weights aren't updated
+		model = SMTModelForCausalLM.from_pretrained("antoniorv6/smt-grandstaff")
+		model.to(device=model.positional_2D.pe.device)
+		input_image = np.mean(input_image, axis=2, keepdims=True) # 3 channels to one
+		input_image = np.transpose(input_image, (2,0,1))[None, :] # add batch size as well, [B, C, H, W]
+		input_image = torch.from_numpy(input_image)#.to(device=model.positional_2D.pe.device)
+
+
+	# take from checkpoint variable
+	elif input_type == 1:
+		model = SMT_Trainer.load_from_checkpoint(checkpoint).model
+		model.to(device=model.pos2D.pe.device)
+		input_image = np.mean(input_image, axis=2, keepdims=True) # 3 channels to one
+		input_image = np.transpose(input_image, (2,0,1))[None, :] # add batch size as well, [B, C, H, W]
+		input_image = torch.from_numpy(input_image).to(device=model.pos2D.pe.device)
+
+	input_image = input_image.to(torch.float32)
+
+	# width / height
+	aspect_ratio = input_image.shape[3]/input_image.shape[2]
+
+	# 8 attention layers * [channels | seq_len | extracted_features]
+	#   extracted features is FLAT input_image shape divided by 16
+	predicted_seq, predictions = model.predict(input_image, return_weights=True)
+
+	# seq_len | reduced_h * reduced_w
+	CA_layers = [ ca_layer.mean(dim=1).squeeze() for ca_layer in predictions.cross_attentions ]
+	
+	seq_len = CA_layers[0].shape[0]
+	att_w = round(sqrt(CA_layers[0].shape[1] * aspect_ratio))
+	att_h = round(sqrt(CA_layers[0].shape[1] / aspect_ratio))
+
+	# make the attention 2-D
+	CA_layers = [ att.reshape( seq_len, att_h, att_w ) for att in CA_layers ]
+
+	# convert to numpy
+	CA_layers = [ att.cpu().detach().numpy() for att in CA_layers ]
+	# ^^^ we store this, then generate the actual images to display ONLY whenever the token slider is moved
+
+	overall = np.stack(CA_layers).sum(axis=0)
+
+	## normalize
+	overall_max_value = np.max(overall)
+	if overall_max_value > 1.0:
+		overall /= np.max(overall)
+
+	CA_layers.append(overall)
+
+	return pd.DataFrame([predicted_seq])
+
+def define_input_source( choice:gr.SelectData ):
+	"""
+	Defines the interface according to the inputs the user has chosen to work with
+	"""
+
+	if choice.index == 0: 		# pretrained weights
+		return 	gr.update(visible=False), 0	 	# file input invisible, input type state update
+
+	elif choice.index == 1: 	# your own weights
+		return 	gr.update(visible=True), 1		# file input visible, input type state update
+
+def define_interface():
+
+	# main components	
+	file_input = gr.File(label="Model Checkpoint File", visible=False, interactive=True)
+	image_input = gr.Image(label="Input Image")
+	tabs = gr.Tabs()
+	
+	# knob components
+	token_slider = gr.Slider(minimum=0, maximum=0, step=1, 
+							label="Pick a token", 
+							info="Select a predicted token to visualize the attention it pays in the input sample",
+							visible=False)
+
+	intensifier_slider = gr.Slider(minimum=1, maximum=100, step=1,
+									label="Intensify attention",
+									info="Use this slider to intensify the attention values to better see differences",
+									value = 10,
+									visible=False)
+
+	token_table = gr.DataFrame(interactive=False, value=pd.DataFrame(["The predicted sequence will appear here"]))
+
+	def intensifier_visibility():
+		"""
+		Makes intensifier slider visible whenever token slider is changed
+		"""
+		return gr.update(visible=True)
+
+	with gr.Blocks() as page:
+
+		###
+		token_slider.release( fn=intensifier_visibility, outputs=intensifier_slider )
+		###
+
+		gr.Markdown("# SMT Demonstrator")
+
+		with gr.Row():
+
+			with gr.Column():
+
+				'''
+				model_interface = gr.Interface(make_predictions,
+									inputs=[file_input, image_input, input_type],
+									outputs=[token_table],
+									flagging_mode='never')
+				'''
+
+				# input area
+				with gr.Blocks():
+					select_src_weights = gr.Dropdown(["Test pretrained weights (default)", "Test your own weights"], 
+													label="Pick which weights to test out", 
+													interactive=True)
+
+					# State variable -- Weights source picked by user
+					input_type = gr.Number(value=0, visible=False)
+
+					select_src_weights.select( define_input_source, outputs=[file_input, input_type] )
+
+					file_input.render()
+					image_input.render()
+
+					with gr.Row():
+
+						def submit_logic(file, image, type):
+
+							return 	make_predictions(file, image, type), gr.update(visible=True), gr.update(visible=True)
+
+						clear_btn = gr.ClearButton( components=[file_input, image_input] )
+
+						submit_btn = gr.Button( value="Submit", variant="primary")
+						submit_btn.click( fn=submit_logic,
+										inputs=[file_input, image_input, input_type],
+										outputs=[token_table, token_slider, intensifier_slider] )
+
+			with gr.Column(scale=2):
+
+				token_slider.render()
+
+				# State variable -- Tab the user left off on
+				tab_selected = gr.Number(value="8", visible=False) # on Overall Attention tab by default
+				
+				# genera las imagenes cada vez que se mueve el slider
+				@gr.render( inputs	=[token_table, token_slider, image_input, intensifier_slider, tab_selected],
+							triggers=[token_slider.release, intensifier_slider.release, token_table.change])
+				def render_images_display(prediction, slider, image, intensifier, tab_no):
+
+					if prediction.shape[0] > 0:
+
+						images = generate_CA_images(slider, image, intensifier)
+
+						gr.Markdown(value="## Contents of the Cross-Attention layers")
+
+						with gr.Tabs(selected=f"{tab_no}") as tabs:
+
+							with gr.Tab(f"Overall", id="8") as tab_overall:
+								tab_overall.select( (lambda : gr.Number(8)), outputs=[tab_selected] )
+								gr.Image(value=images[8])
+
+							with gr.Tab(f"Layer 1", id=f"0") as tab_1:
+									tab_1.select( (lambda : gr.Number(0)), outputs=[tab_selected] )
+									gr.Image(value=images[0])
+
+							with gr.Tab(f"Layer 2", id=f"1") as tab_2:
+									tab_2.select( (lambda : gr.Number(1)), outputs=[tab_selected] )
+									gr.Image(value=images[1])
+
+							with gr.Tab(f"Layer 3", id=f"2") as tab_3:
+									tab_3.select( (lambda : gr.Number(2)), outputs=[tab_selected] )
+									gr.Image(value=images[2])
+
+							with gr.Tab(f"Layer 4", id=f"3") as tab_4:
+									tab_4.select( (lambda : gr.Number(3)), outputs=[tab_selected] )
+									gr.Image(value=images[3])
+
+							with gr.Tab(f"Layer 5", id=f"4") as tab_5:
+									tab_5.select( (lambda : gr.Number(4)), outputs=[tab_selected] )
+									gr.Image(value=images[4])
+
+							with gr.Tab(f"Layer 6", id=f"5") as tab_6:
+									tab_6.select( (lambda : gr.Number(5)), outputs=[tab_selected] )
+									gr.Image(value=images[5])
+
+							with gr.Tab(f"Layer 7", id=f"6") as tab_7:
+									tab_7.select( (lambda : gr.Number(6)), outputs=[tab_selected] )
+									gr.Image(value=images[6])
+
+							with gr.Tab(f"Layer 8", id=f"7") as tab_8:
+									tab_8.select( (lambda : gr.Number(7)), outputs=[tab_selected] )
+									gr.Image(value=images[7])
+
+				intensifier_slider.render()
+
+		with gr.Column():
+
+
+			gr.Markdown("## Predicted Sequence")
+
+			def render_prediction_display(tokens):
+				return gr.Slider(maximum=tokens.shape[0], visible=True), gr.update(visible=True)
+
+			token_table.render()
+			token_table.change(render_prediction_display, inputs=[token_table], outputs=[token_slider, token_table])
+
+
+
+	return page
+
+if __name__=="__main__":
+	page = define_interface()
+	page.launch(share=False)
+