make a program more efficient

app.py

@@ -1,25 +1,25 @@
-import gradio as gr
-from crowd_counter import CrowdCounter
-from PIL import Image
-import cv2
-
-crowd_counter_engine = CrowdCounter()
-
-def predict(inp):
-    inp = Image.fromarray(inp.astype('uint8'), 'RGB')
-    response = crowd_counter_engine.inference(inp)
-    crowd_count = response[0]
-    pred_img= response[1]
-    return cv2.cvtColor(pred_img, cv2.COLOR_BGR2RGB), crowd_count
-
-title = "Crowd Counter Demo"
-desc = "A Demo of Proposal Point Prediction for Crowd Counting - Powered by P2PNet"
-examples = [
-    ["images/img-1.jpg"],
-    ["images/img-2.jpg"],
-    ["images/img-3.jpg"],
-]
-inputs = gr.inputs.Image(label="Image of Crowd")
-outputs = [gr.outputs.Image(label="Proposal Points Prediction",type = "numpy"), gr.outputs.Label(label="Predicted Count",type = "numpy")]
-gr.Interface(fn=predict, inputs=inputs, outputs=outputs, title=title, description=desc, examples=examples,
-             allow_flagging=False).launch()
+import gradio as gr
+from crowd_counter import CrowdCounter
+from PIL import Image
+import cv2
+
+crowd_counter_engine = CrowdCounter()
+
+def predict(inp):
+    inp = Image.fromarray(inp.astype('uint8'), 'RGB')
+    response = crowd_counter_engine.inference(inp)
+    crowd_count = response[0]
+    pred_img= response[1]
+    return cv2.cvtColor(pred_img, cv2.COLOR_BGR2RGB), crowd_count
+
+title = "Crowd Counter Demo"
+desc = "A Demo of Proposal Point Prediction for Crowd Counting - Powered by P2PNet"
+examples = [
+    ["images/img-1.jpg"],
+    ["images/img-2.jpg"],
+    ["images/img-3.jpg"],
+]
+inputs = gr.inputs.Image(label="Image of Crowd")
+outputs = [gr.outputs.Image(label="Proposal Points Prediction",type = "numpy"), gr.outputs.Label(label="Predicted Count",type = "numpy")]
+gr.Interface(fn=predict, inputs=inputs, outputs=outputs, title=title, description=desc, examples=examples,
+             allow_flagging=False).launch(share=True)

crowd_counter/__init__.py

@@ -40,7 +40,6 @@ class CrowdCounter:
             line=2,
             output_dir="./crowd_counter/preds",
             weight_path="./crowd_counter/weights/SHTechA.pth",
-            # gpu_id=0,
         )
 
         # device = torch.device('cuda')
@@ -68,14 +67,25 @@ class CrowdCounter:
     def test(
         self, args: Args, img_raw: Image.Image , debug: bool = False,
     ) -> tuple[any, Image.Image, torch.Tensor]:
-        
-        # round the size
-        width, height = img_raw.size
-        new_width = width // 128 * 128
-        new_height = height // 128 * 128
-        img_raw = img_raw.resize((new_width, new_height), Image.LANCZOS)
+
+        ori_width, ori_height = img_raw.size
+        max_dimension = 384
+
+        if ori_width > max_dimension or ori_height > max_dimension:
+            scale_factor = max_dimension / max(ori_width, ori_height)
+            new_width = int(ori_width * scale_factor)
+            new_height = int(ori_height * scale_factor)
+            new_width = new_width // 128 * 128
+            new_height = new_height // 128 * 128
+            img_resized = img_raw.resize((new_width, new_height), Image.LANCZOS)
+        else:
+            img_resized = img_raw
+            new_width = ori_width
+            new_height = ori_height
+
+        print(new_width, new_height)
         # pre-proccessing
-        img = self.transform(img_raw)
+        img = self.transform(img_resized)
 
         samples = torch.Tensor(img).unsqueeze(0)
         samples = samples.to(self.device)
@@ -94,15 +104,25 @@ class CrowdCounter:
             outputs_points[outputs_scores > threshold].detach().cpu().numpy().tolist()
         )
 
+        scale_factor_width = ori_width / new_width
+        scale_factor_height = ori_height / new_height
+
+        adjusted_points = []
+        for p in points:
+            # Adjust each point's coordinates
+            adjusted_x = int(p[0] * scale_factor_width)
+            adjusted_y = int(p[1] * scale_factor_height)
+            adjusted_points.append((adjusted_x, adjusted_y))
+
         # draw the predictions
-        size = 5
+        size = 3
         img_to_draw = cv2.cvtColor(np.array(img_raw), cv2.COLOR_RGB2BGR)
 
-        for p in points:
+        for p in adjusted_points:
             img_to_draw = cv2.circle(
-                img_to_draw, (int(p[0]), int(p[1])), size, (255, 0, 0), -1
+                img_to_draw, (int(p[0]), int(p[1])), size, (0, 0, 255), -1
             )
-        return points, img_to_draw, conf
+        return adjusted_points, img_to_draw, conf
 
     # Function to process and save images
     def inference(self, img_raw: Image.Image) -> tuple[int, Image.Image]: