Upload 2 files

app.py

@@ -69,6 +69,7 @@ tracker = Tracker(
     stride=8,
     inference_iters=4,
     target_res=1024,
+    device=device,
 )
 
 # -------------------- Step 1: Extract the First Frame -------------------- #

demo_dense_visualize.py

@@ -1,9 +1,17 @@
 import os
 import random
 import torch
+import signal
+import socket
 import sys
+import json
 import torch.nn.functional as F
 import numpy as np
+import argparse
+from pathlib import Path
+import torch.optim as optim
+from torch.cuda.amp import GradScaler
+from lightning_fabric import Fabric 
 
 import utils.loss
 import utils.samp
@@ -11,10 +19,16 @@ import utils.data
 import utils.improc
 import utils.misc
 import utils.saveload
+from tensorboardX import SummaryWriter
+import datetime
+import time
 import cv2
 import imageio
 from nets.blocks import InputPadder
+from tqdm import tqdm
+# from pytorch_lightning.callbacks import BaseFinetuning
 from utils.visualizer import Visualizer
+from torchvision.transforms.functional import resize
 
 import torch
 import requests
@@ -43,6 +57,42 @@ def run_example(processor, model, task_prompt, image, text_input=None):
 
     return parsed_answer
 
+
+def polygons_to_mask(image, prediction, fill_value=255):
+    """
+    Converts polygons into a mask.
+    
+    Parameters:
+    - image: A PIL Image instance whose size will be used for the mask.
+    - prediction: Dictionary containing 'polygons' and 'labels'. 
+                  'polygons' is a list where each element is a list of sub-polygons.
+    - fill_value: The pixel value used to fill the polygon areas (default 255 for a binary mask).
+    
+    Returns:
+    - A NumPy array representing the mask (same width and height as the input image).
+    """
+    # Create a blank grayscale mask image with the same size as the original image.
+    mask = Image.new('L', image.size, 0)
+    draw = ImageDraw.Draw(mask)
+    
+    # Iterate over each set of polygons
+    for polygons in prediction['polygons']:
+        # Each element in "polygons" can be a sub-polygon
+        for poly in polygons:
+            # Ensure the polygon is in the right shape and has at least 3 points.
+            poly_arr = np.array(poly).reshape(-1, 2)
+            if poly_arr.shape[0] < 3:
+                print('Skipping invalid polygon:', poly_arr)
+                continue
+            # Convert the polygon vertices into a list for drawing.
+            poly_list = poly_arr.reshape(-1).tolist()
+            # Draw the polygon on the mask with the fill_value.
+            draw.polygon(poly_list, fill=fill_value)
+    
+    # Convert the PIL mask image to a NumPy array and return it.
+    return np.array(mask)
+
+
 class Tracker:
     def __init__(self, model, mean, std, S, stride, inference_iters, target_res, device='cuda'):
         """