Add backbone choice

app.py

@@ -328,13 +328,19 @@ def get_pose_estimation_prediction(pose_model, image, center, scale):
         return preds
 
 
-def main(image_bgr, box_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)):
+def main(image_bgr, backbone_choice, box_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)):
     CTX = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
     
     box_model.to(CTX)
     box_model.eval()
-    model = torch.hub.load('yangsenius/TransPose:main', 'tph_a4_256x192', pretrained=True)
+    
+    if backbone_choice == "HRNet":
+      backbone_choice = "tph_a4_256x192"
+    else:
+      backbone_choice == "ResNet"
+      backbone_choice = "tpr_a4_256x192"
+    model = torch.hub.load('yangsenius/TransPose:main', backbone_choice , pretrained=True)
 
     img_dimensions = (256, 192)
 
@@ -355,8 +361,7 @@ def main(image_bgr, box_model = torchvision.models.detection.fasterrcnn_resnet50
                 for kpt in pose_preds:
                     draw_pose(kpt, image_bgr)  # draw the poses
 
-    im = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
-    return im
+    return image_bgr
 
 title = "TransPose"
 description = "Gradio demo for TransPose: Keypoint localization via Transformer. Dataset: COCO train2017 & COCO val2017."
@@ -364,5 +369,5 @@ article = "<div style='text-align: center;'><a href='https://github.com/yangseni
 
 examples = [["./examples/one.jpg"], ["./examples/two.jpg"]]
 
-iface = gr.Interface(main, inputs=gr.inputs.Image(), outputs="image", description=description, article=article, title=title, examples=examples)
+iface = gr.Interface(main, inputs=[gr.inputs.Image(), gr.inputs.Radio(["HRNet", "ResNet"])], outputs="image", description=description, article=article, title=title, examples=examples)
 iface.launch(enable_queue=True, debug='True')