Spaces:
Running
Running
| import torch | |
| import gradio as gr | |
| import json | |
| from torchvision import transforms | |
| from PIL import Image, ImageDraw, ImageFont | |
| TORCHSCRIPT_PATH = "res/screenrecognition-web350k-vins.torchscript" | |
| LABELS_PATH = "res/class_map_vins_manual.json" | |
| model = torch.jit.load(TORCHSCRIPT_PATH) | |
| with open(LABELS_PATH, "r") as f: | |
| idx2Label = json.load(f)["idx2Label"] | |
| img_transforms = transforms.ToTensor() | |
| # inter_class_nms and iou functions implemented by GPT | |
| def inter_class_nms(boxes, scores, iou_threshold=0.5): | |
| # Convert boxes and scores to torch tensors if they are not already | |
| boxes = torch.as_tensor(boxes) | |
| scores, class_indices = scores.max(dim=1) | |
| # Keep track of final boxes and scores | |
| final_boxes = [] | |
| final_scores = [] | |
| final_class_indices = [] | |
| for class_index in range(scores.shape[1]): | |
| # Filter boxes and scores for the current class | |
| class_scores = scores[:, class_index] | |
| class_boxes = boxes | |
| # Indices of boxes sorted by score (highest first) | |
| sorted_indices = torch.argsort(class_scores, descending=True) | |
| while len(sorted_indices) > 0: | |
| # Take the box with the highest score | |
| highest_index = sorted_indices[0] | |
| highest_box = class_boxes[highest_index] | |
| # Add the highest box and score to the final list | |
| final_boxes.append(highest_box) | |
| final_scores.append(class_scores[highest_index]) | |
| final_class_indices.append(class_index) | |
| # Remove the highest box from the list | |
| sorted_indices = sorted_indices[1:] | |
| # Compute IoU of the highest box with the rest | |
| ious = iou(class_boxes[sorted_indices], highest_box) | |
| # Keep only boxes with IoU less than the threshold | |
| sorted_indices = sorted_indices[ious < iou_threshold] | |
| return {'boxes': final_boxes, 'scores': final_scores} | |
| def iou(boxes1, boxes2): | |
| """ | |
| Compute the Intersection over Union (IoU) of two sets of boxes. | |
| Args: | |
| - boxes1 (Tensor[N, 4]): ground truth boxes | |
| - boxes2 (Tensor[M, 4]): predicted boxes | |
| Returns: | |
| - iou (Tensor[N, M]): the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2 | |
| """ | |
| area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1]) | |
| area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1]) | |
| lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2] | |
| rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2] | |
| wh = (rb - lt).clamp(min=0) # [N,M,2] | |
| inter = wh[:, :, 0] * wh[:, :, 1] # [N,M] | |
| iou = inter / (area1[:, None] + area2 - inter) | |
| return iou | |
| def predict(img, conf_thresh=0.4): | |
| img_input = [img_transforms(img)] | |
| _, pred = model(img_input) | |
| pred = inter_class_nms(pred['boxes'], pred['scores']) | |
| out_img = img.copy() | |
| draw = ImageDraw.Draw(out_img) | |
| font = ImageFont.truetype("res/Tuffy_Bold.ttf", 25) | |
| for i in range(len(pred[0]['boxes'])): | |
| conf_score = pred[0]['scores'][i] | |
| if conf_score > conf_thresh: | |
| x1, y1, x2, y2 = pred[0]['boxes'][i] | |
| x1 = int(x1) | |
| y1 = int(y1) | |
| x2 = int(x2) | |
| y2 = int(y2) | |
| draw.rectangle([x1, y1, x2, y2], outline='red', width=3) | |
| text = idx2Label[str(int(pred[0]['labels'][i]))] + " {:.2f}".format(float(conf_score)) | |
| bbox = draw.textbbox((x1, y1), text, font=font) | |
| draw.rectangle(bbox, fill="red") | |
| draw.text((x1, y1), text, font=font, fill="black") | |
| return out_img | |
| example_imgs = [ | |
| ["res/example.jpg", 0.4], | |
| ["res/screenlane-snapchat-profile.jpg", 0.4], | |
| ["res/screenlane-snapchat-settings.jpg", 0.4], | |
| ["res/example_pair1.jpg", 0.4], | |
| ["res/example_pair2.jpg", 0.4], | |
| ] | |
| interface = gr.Interface(fn=predict, inputs=[gr.Image(type="pil", label="Screenshot"), gr.Slider(0.0, 1.0, step=0.1, value=0.4)], outputs=gr.Image(type="pil", label="Annotated Screenshot").style(height=600), examples=example_imgs) | |
| interface.launch() | |