import gradio as gr import cv2 from PIL import Image import timm from ultralytics import YOLO import numpy as np import torch from PIL import Image as PILImage from timm.utils.attention_extract import AttentionExtract from glob import glob import os import argparse from gradio_utils import classes, cropSquareCV, nms, classifier_transforms, get_sns_heatmap from gradio_utils import make_thumbnail_gallery, annotate_bbox_on_image from models import MCAClassifier device = "cuda" if torch.cuda.is_available() else "cpu" dirname = os.path.dirname(__file__) detector = YOLO(os.path.join(dirname, 'models/MDV6-yolov10-e-1280.pt')) classifier = timm.create_model("vit_large_patch14_dinov2.lvd142m", pretrained=False, num_classes=46, dynamic_img_size=True) classifier.load_state_dict(torch.load(os.path.join(dirname, 'models/crop_classifier_half.pt'), map_location=device, weights_only=True)) classifier.to(device) seq_model = MCAClassifier(depth=4) seq_model.load_state_dict(torch.load(os.path.join(dirname, 'models/mca_classifier.pt'), map_location=device, weights_only=True)) seq_model.to(device) extractor = AttentionExtract(seq_model, method='fx') # for the attention matrix def show_preds_image(image1, image2=None): image_paths = [image1, image2] predictions = [] embeddings_seq = [] gallery_images = [] table_data = [] img_list = [] crop_size = (200, 200) # detection and independent classifier + save embeddings for mca classsifer for image_path in image_paths: if image_path is None: continue image = cv2.imread(image_path) outputs = detector.predict(source=image_path, imgsz=1280, verbose=True, conf=0.25, device=device) detections = nms(outputs[0].cpu().numpy().boxes, 0.5) for i, (box_cls, conf, box) in enumerate(zip(detections.cls, detections.conf, detections.xyxy)): if conf < 0.25: continue x1, y1, x2, y2 = map(int, box) croppedimagecv = cropSquareCV(image, box.copy())[:, :, (2, 1, 0)] croppedimage = Image.fromarray(croppedimagecv) # for classifier im = cv2.resize(croppedimagecv, crop_size, interpolation=cv2.INTER_CUBIC) # for visualization gallery_images.append(PILImage.fromarray(im)) # for gallery img_list.append(im) # for attention matrix if box_cls == 0: x = classifier_transforms(croppedimage).unsqueeze(dim=0).to(device) with torch.no_grad(): embeddings = classifier.forward_features(x) embeddings_seq.append(embeddings[:, 0:1, :]) pred = classifier.forward_head(embeddings).softmax(dim=1).cpu() scores = {classes[j]: float(pred[0][j]) for j in range(len(classes))} top_class = max(scores, key=scores.get) label_text = f"{top_class}: {scores[top_class]*100:.0f}%" else: if box_cls == 1: label_text = "human" elif box_cls == 2: label_text = "vehicle" predictions.append([(x1, y1, x2, y2), label_text, image_path]) # sequence-aware classifier if len(embeddings_seq) > 0: x = torch.concat(embeddings_seq, 1) with torch.no_grad(): embeddings = seq_model.forward_blocks(x) pred = classifier.forward_head(embeddings.swapaxes(0, 1)).cpu().softmax(dim=1) oo = extractor(x) am = oo[f'blocks.{len(oo)-1}.attn.softmax'].cpu().numpy()[0] print(am.shape) else: pred = [] am = None n = 0 for i, ((x1, y1, x2, y2), orig_label, image_path) in enumerate(predictions): if orig_label in ["human", "vehicle"]: table_data.append([i, orig_label, orig_label, image_path]) else: top_class = max({classes[j]: float(pred[n][j]) for j in range(len(classes))}, key=lambda k: pred[n][classes.index(k)]) top_score = float(pred[n][classes.index(top_class)])*100 seq_label = f"{top_class}: {top_score:.0f}%" table_data.append([i, orig_label, seq_label, image_path]) n += 1 # permutate bbox order to make the attention matrix easier to read sp = [table_data[i][2].split(":")[0] for i in range(len(table_data))] perm = sorted(range(len(sp)), key=lambda k: sp[k]) gallery_images = [gallery_images[i] for i in perm] img_list = [img_list[i] for i in perm] table_data = [table_data[i] for i in perm] if am is not None: not_animal = [x in ["human", "vehicle"] for x in [sp[i] for i in perm]] not_animal_idx = [perm[i] for i in range(len(perm)) if not_animal[i]] am_perm = [] for i in range(len(perm)): if not_animal[i]: continue else: am_perm.append(perm[i] - sum([perm[i] > k for k in not_animal_idx])) am = am[:, am_perm, :] am = am[:, :, am_perm] img_list = [img_list[i] for i in range(len(img_list)) if not not_animal[i]] # annotate bounding boxes on images annotated_image1 = cv2.imread(image1) if image1 else None annotated_image2 = cv2.imread(image2) if image2 else None for i, ((x1, y1, x2, y2), orig_label, image_path) in enumerate(predictions): if image_path == image1: annotate_bbox_on_image(annotated_image1, x1, y1, x2, y2, str(perm.index(i)+1)) else: annotate_bbox_on_image(annotated_image2, x1, y1, x2, y2, str(perm.index(i)+1)) gallery_images = make_thumbnail_gallery(gallery_images, table_data) # make attention matrix img = np.concatenate(img_list, axis=1) X, Y, _ = img.shape am_img = np.zeros((X+Y, X+Y, 3), dtype=np.uint8) if am is not None: am_img[:X, X:X+Y] = img am_img[X:X+Y, :X] = np.concatenate(img_list, axis=0) am_img[X:, X:] = cv2.resize(get_sns_heatmap(am), (Y, Y), interpolation=cv2.INTER_NEAREST) output_image1 = cv2.cvtColor(annotated_image1, cv2.COLOR_BGR2RGB) if image1 else None output_image2 = cv2.cvtColor(annotated_image2, cv2.COLOR_BGR2RGB) if image2 else None return output_image1, output_image2, gallery_images, am_img with gr.Blocks(css=".block {margin: auto}") as demo: args = dict(show_download_button=False, show_fullscreen_button=False) if False else dict() with gr.Tab("One image"): with gr.Row(): with gr.Column(scale=1): input_image = gr.Image(type="filepath", label="Image") path = [x for x in glob(os.path.join(dirname, 'images/gradio_examples/one_image/*')) if x.lower().split(".")[-1] in ["jpg", "jpeg", "png"]] examples = gr.Examples(path, input_image) btn = gr.Button("Run", variant="primary") with gr.Column(scale=3): with gr.Group(): with gr.Row(): output_image = gr.Image(type="numpy", label="Annotated Image", scale=16, show_label=False, **args) output_image2 = gr.Image(type="numpy", label="Annotated Image 2", scale=0, render=False, show_label=False, **args) am_image = gr.Image(type="numpy", label="Attention Map", scale=9, show_label=False, **args) with gr.Row(): output_gallery = gr.Gallery(label="Detected Crops", columns=9, rows=2, height="auto", object_fit='scale-down', show_label=False) btn.click( fn=show_preds_image, inputs=[input_image], outputs=[output_image, output_image2, output_gallery, am_image] ) with gr.Tab("Two images"): with gr.Row(): with gr.Column(scale=1): input_image = gr.Image(type="filepath", label="Image") path2 = [x for x in glob(os.path.join(dirname, 'images/gradio_examples/two_images/*')) if x.lower().split(".")[-1] in ["jpg", "jpeg", "png"]] path2.sort() examples = gr.Examples(path2, input_image) input_image2 = gr.Image(type="filepath", label="Second Image (optional)") examples2 = gr.Examples(path2, input_image2) btn = gr.Button("Run", variant="primary") with gr.Column(scale=3): with gr.Group(): with gr.Row(): output_image = gr.Image(type="numpy", label="Annotated Image", scale=4, show_label=False, **args) output_image2 = gr.Image(type="numpy", label="Annotated Image 2", scale=4, show_label=False, **args) am_image = gr.Image(type="numpy", label="Attention Map", scale=3, show_label=False, **args) with gr.Row(): output_gallery = gr.Gallery(label="Detected Crops", columns=9, rows=2, height="auto", object_fit='scale-down', show_label=False) btn.click( fn=show_preds_image, inputs=[input_image, input_image2], outputs=[output_image, output_image2, output_gallery, am_image] ) demo.launch(allowed_paths=[os.path.join(dirname, '../images/gradio_examples/')])