Update app.py

app.py

@@ -1,5 +1,77 @@
 import gradio as gr
 import torch
+from io import BytesIO
+
+import cv2
+import gradio as gr
+import numpy as np
+import requests
+from PIL import Image
+
+
+from super_gradients.common.object_names import Models
+from super_gradients.training import models
+from super_gradients.training.utils.visualization.detection import draw_bbox
+
+
+# Initialize your pose estimation model
+yolo_nas_pose = models.get("best.pt",
+                           num_classes=1,
+                           checkpoint_path="./best.pt")
+
+def process_and_predict(url=None,
+                        image=None,
+                        confidence=0.5,
+                        iou=0.5):
+    # If a URL is provided, use it directly for prediction
+    if url is not None and url.strip() != "":
+        response = requests.get(url)
+        image = Image.open(BytesIO(response.content))
+        image = np.array(image)
+        result = yolo_nas_pose.predict(image, conf=confidence,iou=iou)
+    # If a file is uploaded, read it, convert it to a numpy array and use it for prediction
+    elif image is not None:
+        result = yolo_nas_pose.predict(image, conf=confidence,iou=iou)
+    else:
+        return None  # If no input is provided, return None
+
+    # Extract prediction data
+    image_prediction = result._images_prediction_lst[0]
+
+    pose_data = image_prediction.prediction
+
+    # Visualize the prediction
+    output_image = PoseVisualization.draw_poses(
+        image=image_prediction.image,
+        poses=pose_data.poses,
+        boxes=pose_data.bboxes_xyxy,
+        scores=pose_data.scores,
+        is_crowd=None,
+        edge_links=pose_data.edge_links,
+        edge_colors=pose_data.edge_colors,
+        keypoint_colors=pose_data.keypoint_colors,
+        joint_thickness=2,
+        box_thickness=2,
+        keypoint_radius=5
+    )
+
+    blank_image = np.zeros_like(image_prediction.image)
+
+    skeleton_image = PoseVisualization.draw_poses(
+    image=blank_image,
+    poses=pose_data.poses,
+    boxes=pose_data.bboxes_xyxy,
+    scores=pose_data.scores,
+    is_crowd=None,
+    edge_links=pose_data.edge_links,
+    edge_colors=pose_data.edge_colors,
+    keypoint_colors=pose_data.keypoint_colors,
+    joint_thickness=2,
+    box_thickness=2,
+    keypoint_radius=5
+)
+
+    return output_image, skeleton_image
 
 def greet(name):
     return "Hello " + name + "!!"
@@ -15,27 +87,6 @@ examples = [ # need to manually delete cache everytime new examples are added
     ["clip2_-1450-_jpg.jpg"], 
     ["clip2_-539-_jpg.jpg"]]
 
-def speclab(img):
-
-    # initialize the model
-    model = torch.hub.load('SamDaaLamb/ValorantTracker', 'srdetect', force_reload=True) # for some reasons loads the model in src rather than demo
-    model.eval()
-
-    # preprocess image to be used as input
-    transforms = A.Compose([
-        A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
-        ToTensorV2()
-    ])
-    input = transforms(image=img)['image']
-    input = input.unsqueeze(0)
-
-    # model prediction
-    output = model(input)
-
-    # overlay output onto original image
-    img[output==255] = [0, 255, 0]
-
-    return img
 
 # define app features and run
 title = "SpecLab Demo"