Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| # -*- coding: utf-8 -*- | |
| from torch import tensor | |
| from torchvision.transforms import InterpolationMode, transforms | |
| import timm | |
| from PIL import Image | |
| import cv2 | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import seaborn as sns | |
| from io import BytesIO | |
| timm.layers.set_fused_attn(False) | |
| def create_text_rectangle(text, width=200, height=40, font=cv2.FONT_HERSHEY_SIMPLEX, color_bgr=(102, 163, 255), | |
| color_text=(0, 0, 0), text_pad_h=10): | |
| image = np.full((height, width, 3), color_bgr, dtype=np.uint8) | |
| font_scale = 1.0 | |
| thickness = 1 | |
| while True: | |
| (text_width, text_height), _ = cv2.getTextSize(text, font, font_scale, thickness) | |
| if text_width <= width - 5 and text_height <= height - text_pad_h: | |
| break | |
| font_scale -= 0.1 | |
| if font_scale <= 0.1: # Avoid going to zero | |
| font_scale = 0.1 | |
| break | |
| x = (width - text_width) // 2 | |
| y = (height + text_height) // 2 | |
| cv2.putText(image, text, (x, y), font, font_scale, color_text, thickness, cv2.LINE_AA) | |
| return image | |
| def nms(detections, iou_threshold=0.5): | |
| if detections is None or len(detections) == 0: | |
| return detections | |
| boxes = detections.xyxy | |
| scores = detections.conf | |
| x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] | |
| areas = (x2 - x1) * (y2 - y1) | |
| order = scores.argsort()[::-1] | |
| keep = [] | |
| while order.size > 0: | |
| i = order[0] | |
| keep.append(i) | |
| xx1 = np.maximum(x1[i], x1[order[1:]]) | |
| yy1 = np.maximum(y1[i], y1[order[1:]]) | |
| xx2 = np.minimum(x2[i], x2[order[1:]]) | |
| yy2 = np.minimum(y2[i], y2[order[1:]]) | |
| w = np.maximum(0.0, xx2 - xx1) | |
| h = np.maximum(0.0, yy2 - yy1) | |
| inter = w * h | |
| iou = inter / (areas[i] + areas[order[1:]] - inter) | |
| inds = np.where(iou <= iou_threshold)[0] | |
| order = order[inds + 1] | |
| return detections[keep] | |
| convert = { | |
| 'acinonyxjubatus': 'Cheetah', | |
| 'aepycerosmelampus': 'Impala', | |
| 'alcelaphusbuselaphus': 'Hartebeest', | |
| 'ardeotiskori': 'Kori Bustard', | |
| 'aves': 'Bird', | |
| 'caracalcaracal': 'Caracal', | |
| 'chlorocebuspygerythrus': 'Vervet Monkey', | |
| 'connochaetestaurinus': 'Wildebeest', | |
| 'crocutacrocuta': 'Spotted Hyena', | |
| 'damaliscuslunatusjimela': 'Topi', | |
| 'equusquagga': 'Zebra', | |
| 'eudorcasthomsonii': "Thomson's Gazelle", | |
| 'felislybica': 'African Wildcat', | |
| 'genetta': 'Genets', | |
| 'giraffacamelopardalis': 'Giraffe', | |
| 'herpestidae': 'Mongooses', | |
| 'hippopotamusamphibius': 'Hippopotamus', | |
| 'hyaenahyaena': 'Striped Hyena', | |
| 'hystrixcristata': 'Porcupine', | |
| 'ictonyxstriatus': 'Polecat', | |
| 'kobusellipsiprymnus': 'Waterbuck', | |
| 'leptailurusserval': 'Serval', | |
| 'lepus': 'Hares', | |
| 'loxodontaafricana': 'Elephant', | |
| 'lupulellamesomelas': 'Jackal', | |
| 'madoqua': 'Dik-diks', | |
| 'mellivoracapensis': 'Honey Badger', | |
| 'nangergranti': "Grant's Gazelle", | |
| 'numididae': 'Guineafowl', | |
| 'orycteropusafer': 'Aardvark', | |
| 'otocyonmegalotis': 'Bat-eared Fox', | |
| 'pantheraleo': 'Lion', | |
| 'pantherapardus': 'Leopard', | |
| 'papio': 'Baboons', | |
| 'phacochoerusafricanus': 'Warthog', | |
| 'protelescristatus': 'Aardwolf', | |
| 'redunca': 'Reedbucks', | |
| 'reptilia': 'Reptiles', | |
| 'rhinocerotidae': 'Rhinoceroses', | |
| 'rodentia': 'Rodents', | |
| 'sagittariusserpentarius': 'Secretarybird', | |
| 'struthiocamelus': 'Ostrich', | |
| 'synceruscaffer': 'Buffalo', | |
| 'tragelaphusoryx': 'Eland', | |
| 'tragelaphusscriptus': 'Bushbuck', | |
| 'viverridae': 'Viverrids' | |
| } | |
| classes = list(convert.values()) | |
| def cropSquareCV(imagecv, box): | |
| x1, y1, x2, y2 = box | |
| xsize = (x2-x1) | |
| ysize = (y2-y1) | |
| if xsize > ysize: | |
| y1 = y1-int((xsize-ysize)/2) | |
| y2 = y2+int((xsize-ysize)/2) | |
| if ysize > xsize: | |
| x1 = x1-int((ysize-xsize)/2) | |
| x2 = x2+int((ysize-xsize)/2) | |
| height, width, _ = imagecv.shape | |
| croppedimagecv = imagecv[max(0, int(y1)): min(int(y2), height), max(0, int(x1)): min(int(x2), width)] | |
| return croppedimagecv | |
| classifier_transforms = transforms.Compose([ | |
| transforms.Resize(size=(182, 182), interpolation=InterpolationMode.BICUBIC, max_size=None, antialias=None), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=tensor([0.4850, 0.4560, 0.4060]), std=tensor([0.2290, 0.2240, 0.2250]))]) | |
| def get_sns_heatmap(attention_matrix): | |
| d = attention_matrix.mean(0) | |
| d_norm = (d - d.min()) / (d.max() - d.min()) | |
| # Plot the heatmap with seaborn | |
| fig, ax = plt.subplots(figsize=(3, 3), dpi=100) | |
| sns.heatmap(d_norm, cmap=sns.light_palette("#825f87", as_cmap=True), cbar=False, xticklabels=False, yticklabels=False, ax=ax) | |
| ax.axis('off') | |
| # Save the heatmap to a buffer | |
| buf = BytesIO() | |
| plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0) | |
| plt.close(fig) | |
| # Load the image from the buffer | |
| buf.seek(0) | |
| heatmap_img = Image.open(buf).convert("RGB") | |
| heatmap_np = np.array(heatmap_img) | |
| return heatmap_np | |
| def make_thumbnail_gallery(gallery_images, table_data): | |
| new_gallery = [] | |
| for i in range(len(gallery_images)): | |
| im = np.concatenate((np.array(gallery_images[i]), create_text_rectangle(table_data[i][1]), | |
| create_text_rectangle(table_data[i][2], color_bgr=(176, 224, 143)))) | |
| im[:30, :30, :] = create_text_rectangle(str(i+1), 30, 30, color_bgr=(0, 0, 0), | |
| color_text=(255, 255, 255), text_pad_h=5) | |
| new_gallery.append(Image.fromarray(im)) | |
| return new_gallery | |
| def annotate_bbox_on_image(annotated_image, x1, y1, x2, y2, label_text=None): | |
| height, width, _ = annotated_image.shape | |
| cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color=(0, 0, 255), thickness=2, lineType=cv2.LINE_AA) | |
| if label_text: | |
| fontScale = fontScale = max(0.5, min(width, height) / 2000) | |
| thickness = 3 | |
| (text_width, text_height), baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, fontScale, thickness) | |
| bg_x1, bg_y1 = x1, max(0, y1 - text_height - baseline - 4) | |
| bg_x2, bg_y2 = x1 + text_width + 4, y1 | |
| cv2.rectangle(annotated_image, (bg_x1, bg_y1), (bg_x2, bg_y2), (0, 0, 0), -1) | |
| cv2.putText(annotated_image, label_text, (x1 + 2, y1 - baseline - 2), cv2.FONT_HERSHEY_SIMPLEX, | |
| fontScale=fontScale, color=(255, 255, 255), thickness=thickness, lineType=cv2.LINE_AA) | |