Update app.py

app.py

@@ -1,42 +1,60 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 import torch
 import gradio as gr
 from gradio_image_prompter import ImagePrompter
-from torch.nn import DataParallel
 from models.counter_infer import build_model
 from utils.arg_parser import get_argparser
 from utils.data import resize_and_pad
 import torchvision.ops as ops
 from torchvision import transforms as T
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image, ImageDraw
 import numpy as np
 
-# Load model (once, to avoid reloading)
+# -----------------------
+# LOAD MODEL
+# -----------------------
 def load_model():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    args = get_argparser().parse_args()
+
+    args = get_argparser().parse_args([])
     args.zero_shot = True
-    model = DataParallel(build_model(args).to(device))
-    model.load_state_dict(torch.load('CNTQG_multitrain_ca44.pth', weights_only=True)['model'], strict=False)
+
+    model = build_model(args).to(device)
+
+    # ⚠️ Make sure this file exists in your repo
+    checkpoint = torch.load("CNTQG_multitrain_ca44.pth", map_location=device)
+
+    model.load_state_dict(checkpoint["model"], strict=False)
     model.eval()
+
     return model, device
 
 model, device = load_model()
 
-# **Function to Process Image Once**
+# -----------------------
+# PROCESS IMAGE
+# -----------------------
 def process_image_once(inputs, enable_mask):
-    model.module.return_masks = enable_mask
+    model.return_masks = enable_mask  # ✅ FIXED
 
     image = inputs['image']
     drawn_boxes = inputs['points']
+
     image_tensor = torch.tensor(image).to(device)
     image_tensor = image_tensor.permute(2, 0, 1).float() / 255.0
-    image_tensor = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(image_tensor)
 
-    bboxes_tensor = torch.tensor([[box[0], box[1], box[3], box[4]] for box in drawn_boxes], dtype=torch.float32).to(device)
+    image_tensor = T.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )(image_tensor)
+
+    bboxes_tensor = torch.tensor(
+        [[box[0], box[1], box[3], box[4]] for box in drawn_boxes],
+        dtype=torch.float32
+    ).to(device)
 
     img, bboxes, scale = resize_and_pad(image_tensor, bboxes_tensor, size=1024.0)
+
     img = img.unsqueeze(0).to(device)
     bboxes = bboxes.unsqueeze(0).to(device)
 
@@ -45,69 +63,44 @@ def process_image_once(inputs, enable_mask):
 
     return image, outputs, masks, img, scale, drawn_boxes
 
-# **Post-process and Update Output**
+# -----------------------
+# POST PROCESS
+# -----------------------
 def post_process(image, outputs, masks, img, scale, drawn_boxes, enable_mask, threshold):
     idx = 0
-    threshold = 1/threshold
-    keep = ops.nms(outputs[idx]['pred_boxes'][outputs[idx]['box_v'] > outputs[idx]['box_v'].max() / threshold],
-                   outputs[idx]['box_v'][outputs[idx]['box_v'] > outputs[idx]['box_v'].max() / threshold], 0.5)
+    threshold = 1 / threshold
+
+    scores = outputs[idx]['box_v']
+    boxes = outputs[idx]['pred_boxes']
 
-    pred_boxes = outputs[idx]['pred_boxes'][outputs[idx]['box_v'] > outputs[idx]['box_v'].max() / threshold][keep]
+    keep_mask = scores > scores.max() / threshold
+
+    keep = ops.nms(
+        boxes[keep_mask],
+        scores[keep_mask],
+        0.5
+    )
+
+    pred_boxes = boxes[keep_mask][keep]
     pred_boxes = torch.clamp(pred_boxes, 0, 1)
 
     pred_boxes = (pred_boxes.cpu() / scale * img.shape[-1]).tolist()
 
-    image = Image.fromarray((image).astype(np.uint8))
-
-    if enable_mask:
-        from matplotlib import pyplot as plt
-        masks_ = masks[idx][(outputs[idx]['box_v'] > outputs[idx]['box_v'].max() / threshold)[0]]
-        N_masks = masks_.shape[0]
-        indices = torch.randint(1, N_masks + 1, (1, N_masks), device=masks_.device).view(-1, 1, 1)
-        masks = (masks_ * indices).sum(dim=0)
-        mask_display = (
-            T.Resize((int(img.shape[2] / scale), int(img.shape[3] / scale)), interpolation=T.InterpolationMode.NEAREST)(
-                masks.cpu().unsqueeze(0))[0])[:image.size[1], :image.size[0]]
-        cmap = plt.cm.tab20
-        norm = plt.Normalize(vmin=0, vmax=N_masks)
-        del masks
-        del masks_
-        del outputs
-        rgba_image = cmap(norm(mask_display))
-        rgba_image[mask_display == 0, -1] = 0
-        rgba_image[mask_display != 0, -1] = 0.5
-
-        overlay = Image.fromarray((rgba_image * 255).astype(np.uint8), mode="RGBA")
-        image = image.convert("RGBA")
-        image = Image.alpha_composite(image, overlay)
-
+    image = Image.fromarray(image.astype(np.uint8))
 
     draw = ImageDraw.Draw(image)
+
     for box in pred_boxes:
         draw.rectangle([box[0], box[1], box[2], box[3]], outline="orange", width=5)
-    # for box in drawn_boxes:
-    #     draw.rectangle([box[0], box[1], box[3], box[4]], outline="red", width=3)
-
-    width, height = image.size
-    square_size = int(0.05 * width)
-    x1, y1 = 10, height - square_size - 10
-    x2, y2 = x1 + square_size, y1 + square_size
-
-    # draw.rectangle([x1, y1, x2, y2], outline="black", fill="black", width=1)
-    # font = ImageFont.load_default()
-    # txt = str(len(pred_boxes))
-    # w = draw.textlength(txt, font=font)
-    # text_x = x1 + (square_size - w) / 2
-    # text_y = y1 + (square_size - 10) / 2
-    # draw.text((text_x, text_y), txt, fill="white", font=font)
 
     return image, len(pred_boxes)
 
-
+# -----------------------
+# GRADIO UI
+# -----------------------
 iface = gr.Blocks()
 
 with iface:
-    # Store intermediate states
     image_input = gr.State()
     outputs_state = gr.State()
     masks_state = gr.State()
@@ -115,45 +108,34 @@ with iface:
     scale_state = gr.State()
     drawn_boxes_state = gr.State()
 
-    # UI Layout: Input Section
     with gr.Row():
         image_prompter = ImagePrompter()
         image_output = gr.Image(type="pil")
-        
 
-    # UI Layout: Output Section
     with gr.Row():
         count_output = gr.Number(label="Total Count")
-        enable_mask = gr.Checkbox(label="Predict masks", value=True)  # Mask enabled by default
-        threshold = gr.Slider(0.05, 0.95, value=0.33, step=0.01, label="Threshold")  # Updated range and default
-        
+        enable_mask = gr.Checkbox(label="Predict masks", value=True)
+        threshold = gr.Slider(0.05, 0.95, value=0.33, step=0.01)
 
-    # Create the 'Count' button
     count_button = gr.Button("Count")
 
-    # Process image once when "Count" button is pressed
     def initial_process(inputs, enable_mask, threshold):
-        # Perform inference once
         image, outputs, masks, img, scale, drawn_boxes = process_image_once(inputs, enable_mask)
 
-        # Save intermediate states
         return (
-            *post_process(image, outputs, masks, img, scale, drawn_boxes, enable_mask, threshold),  # Processed outputs
-            image, outputs, masks, img, scale, drawn_boxes  # Store in states for later use
+            *post_process(image, outputs, masks, img, scale, drawn_boxes, enable_mask, threshold),
+            image, outputs, masks, img, scale, drawn_boxes
         )
 
-    # Update image and count when the threshold slider changes (post-process only)
     def update_threshold(threshold, image, outputs, masks, img, scale, drawn_boxes, enable_mask):
         return post_process(image, outputs, masks, img, scale, drawn_boxes, enable_mask, threshold)
 
-    # Run initial inference and post-process when "Count" button is clicked
     count_button.click(
         initial_process,
-        [image_prompter, enable_mask, threshold],  # Inputs
-        [image_output, count_output, image_input, outputs_state, masks_state, img_state, scale_state, drawn_boxes_state]  # Outputs + States
+        [image_prompter, enable_mask, threshold],
+        [image_output, count_output, image_input, outputs_state, masks_state, img_state, scale_state, drawn_boxes_state]
     )
 
-    # Adjust the output dynamically based on the threshold slider (no re-inference)
     threshold.change(
         update_threshold,
         [threshold, image_input, outputs_state, masks_state, img_state, scale_state, drawn_boxes_state, enable_mask],
@@ -161,10 +143,9 @@ with iface:
     )
 
     enable_mask.change(
-         update_threshold,
+        update_threshold,
         [threshold, image_input, outputs_state, masks_state, img_state, scale_state, drawn_boxes_state, enable_mask],
         [image_output, count_output]
     )
 
-iface.launch(share=True)
-
+iface.launch()