add faster rcnn visualization and avoid reloading model parameters

.gitignore

@@ -2,6 +2,7 @@ gradio_demo_old
 gradio_demo/*.png
 gradio_demo/*.glb
 gradio_cached_examples/
+datasets/test_image_crops_old
 results/gradio_examples/*.png
 results/gradio_examples/*.jpg
 results/gradio_examples/*.glb

README.md

@@ -5,7 +5,7 @@ colorFrom: pink
 colorTo: green
 sdk: gradio
 sdk_version: 3.0.2
-app_file: ./gradio_demo/barc_demo_v3.py
+app_file: ./gradio_demo/barc_demo_v6.py
 pinned: false
 python_version: 3.7.6
 ---

gradio_demo/barc_demo_v6.py

@@ -0,0 +1,291 @@
+# python gradio_demo/barc_demo_v6.py
+
+import os
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"  
+os.environ["CUDA_VISIBLE_DEVICES"]="0"
+try:
+    # os.system("pip install --upgrade  torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html")
+    os.system("pip install --upgrade  torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/cu101/torch_stable.html")
+except Exception as e:
+    print(e)
+
+import numpy as np
+import os
+import glob
+import torch
+from torch.utils.data import DataLoader
+import torchvision
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+import torchvision.transforms as T
+import cv2
+from matplotlib import pyplot as plt
+from PIL import Image
+import random
+import gradio as gr
+
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../', 'src'))
+from stacked_hourglass.datasets.imgcropslist import ImgCrops
+from combined_model.train_main_image_to_3d_withbreedrel import do_visual_epoch
+from combined_model.model_shape_v7 import ModelImageTo3d_withshape_withproj
+from configs.barc_cfg_defaults import get_cfg_global_updated
+
+random.seed(0)
+
+print(
+    "torch: ", torch.__version__,
+    "\ntorchvision: ", torchvision.__version__,
+)
+
+
+
+def get_prediction(model, img_path_or_img, confidence=0.5):
+    """
+    see https://haochen23.github.io/2020/04/object-detection-faster-rcnn.html#.YsMCm4TP3-g
+    get_prediction
+        parameters:
+        - img_path - path of the input image
+        - confidence - threshold value for prediction score
+        method:
+        - Image is obtained from the image path
+        - the image is converted to image tensor using PyTorch's Transforms
+        - image is passed through the model to get the predictions
+        - class, box coordinates are obtained, but only prediction score > threshold
+            are chosen.
+    """
+    if isinstance(img_path_or_img, str):
+        img = Image.open(img_path_or_img).convert('RGB')
+    else:
+        img = img_path_or_img
+    transform = T.Compose([T.ToTensor()])
+    img = transform(img)
+    pred = model([img])
+    # pred_class = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].numpy())]
+    pred_class = list(pred[0]['labels'].numpy())
+    pred_boxes = [[(int(i[0]), int(i[1])), (int(i[2]), int(i[3]))] for i in list(pred[0]['boxes'].detach().numpy())]
+    pred_score = list(pred[0]['scores'].detach().numpy())
+    try:
+        pred_t = [pred_score.index(x) for x in pred_score if x>confidence][-1]
+        pred_boxes = pred_boxes[:pred_t+1]
+        pred_class = pred_class[:pred_t+1]
+        return pred_boxes, pred_class, pred_score
+    except:
+        print('no bounding box with a score that is high enough found! -> work on full image')
+        return None, None, None
+
+
+def detect_object(model, img_path_or_img, confidence=0.5, rect_th=2, text_size=0.5, text_th=1):
+    """
+    see https://haochen23.github.io/2020/04/object-detection-faster-rcnn.html#.YsMCm4TP3-g
+    object_detection_api
+        parameters:
+        - img_path_or_img - path of the input image
+        - confidence - threshold value for prediction score
+        - rect_th - thickness of bounding box
+        - text_size - size of the class label text
+        - text_th - thichness of the text
+        method:
+        - prediction is obtained from get_prediction method
+        - for each prediction, bounding box is drawn and text is written 
+            with opencv
+        - the final image is displayed
+    """
+    boxes, pred_cls, pred_scores = get_prediction(model, img_path_or_img, confidence)
+    if isinstance(img_path_or_img, str):
+        img = cv2.imread(img_path_or_img)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    else:
+        img = img_path_or_img
+    is_first = True
+    bbox = None
+    if boxes is not None:
+        for i in range(len(boxes)):
+            cls = pred_cls[i]
+            if cls == 18 and bbox is None:
+                cv2.rectangle(img, boxes[i][0], boxes[i][1],color=(0, 255, 0), thickness=rect_th)
+                # cv2.putText(img, pred_cls[i], boxes[i][0], cv2.FONT_HERSHEY_SIMPLEX, text_size, (0,255,0),thickness=text_th)
+                # cv2.putText(img, str(pred_scores[i]), boxes[i][0], cv2.FONT_HERSHEY_SIMPLEX, text_size, (0,255,0),thickness=text_th)
+                bbox = boxes[i]
+    return img, bbox
+
+
+# -------------------------------------------------------------------------------------------------------------------- #
+model_bbox = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
+model_bbox.eval()
+
+def run_bbox_inference(input_image):
+    # load configs
+    cfg = get_cfg_global_updated()
+    out_path = os.path.join(cfg.paths.ROOT_OUT_PATH, 'gradio_examples', 'test2.png')
+    img, bbox = detect_object(model=model_bbox, img_path_or_img=input_image, confidence=0.5)
+    fig = plt.figure()   #  plt.figure(figsize=(20,30))
+    plt.imsave(out_path, img)
+    return img, bbox
+# -------------------------------------------------------------------------------------------------------------------- #
+
+
+# -------------------------------------------------------------------------------------------------------------------- #
+# load configs
+cfg = get_cfg_global_updated()
+# Select the hardware device to use for inference.
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print('---> device: ' + device)
+# disable gradient calculations.
+torch.set_grad_enabled(False)
+# prepare complete model
+complete_model = ModelImageTo3d_withshape_withproj(
+    num_stage_comb=cfg.params.NUM_STAGE_COMB, num_stage_heads=cfg.params.NUM_STAGE_HEADS, \
+    num_stage_heads_pose=cfg.params.NUM_STAGE_HEADS_POSE, trans_sep=cfg.params.TRANS_SEP, \
+    arch=cfg.params.ARCH, n_joints=cfg.params.N_JOINTS, n_classes=cfg.params.N_CLASSES, \
+    n_keyp=cfg.params.N_KEYP, n_bones=cfg.params.N_BONES, n_betas=cfg.params.N_BETAS, n_betas_limbs=cfg.params.N_BETAS_LIMBS, \
+    n_breeds=cfg.params.N_BREEDS, n_z=cfg.params.N_Z, image_size=cfg.params.IMG_SIZE, \
+    silh_no_tail=cfg.params.SILH_NO_TAIL, thr_keyp_sc=cfg.params.KP_THRESHOLD, add_z_to_3d_input=cfg.params.ADD_Z_TO_3D_INPUT,
+    n_segbps=cfg.params.N_SEGBPS, add_segbps_to_3d_input=cfg.params.ADD_SEGBPS_TO_3D_INPUT, add_partseg=cfg.params.ADD_PARTSEG, n_partseg=cfg.params.N_PARTSEG, \
+    fix_flength=cfg.params.FIX_FLENGTH, structure_z_to_betas=cfg.params.STRUCTURE_Z_TO_B, structure_pose_net=cfg.params.STRUCTURE_POSE_NET,
+    nf_version=cfg.params.NF_VERSION) 
+# load trained model
+path_model_file_complete = os.path.join(cfg.paths.ROOT_CHECKPOINT_PATH,  'barc_complete', 'model_best.pth.tar') 
+print(path_model_file_complete)
+assert os.path.isfile(path_model_file_complete)
+print('Loading model weights from file: {}'.format(path_model_file_complete))
+checkpoint_complete = torch.load(path_model_file_complete, map_location=device)
+state_dict_complete = checkpoint_complete['state_dict']
+complete_model.load_state_dict(state_dict_complete, strict=False)        
+complete_model = complete_model.to(device)
+# create path for output files
+save_imgs_path = os.path.join(cfg.paths.ROOT_OUT_PATH, 'gradio_examples')
+if not os.path.exists(save_imgs_path):
+    os.makedirs(save_imgs_path)
+
+def run_barc_inference(input_image, bbox=None):
+    input_image_list = [input_image]
+    if bbox is not None:
+        input_bbox_list = [bbox]
+    else:
+        input_bbox_list = None
+    # prepare data loader
+    val_dataset = ImgCrops(image_list=input_image_list, bbox_list=input_bbox_list, dataset_mode='complete')
+    test_name_list = val_dataset.test_name_list
+    val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False,
+                            num_workers=0, pin_memory=True, drop_last=False)   
+    # run visual evaluation
+    all_results = do_visual_epoch(val_loader, complete_model, device,
+                        ImgCrops.DATA_INFO,
+                        weight_dict=None,
+                        acc_joints=ImgCrops.ACC_JOINTS,
+                        save_imgs_path=None, # save_imgs_path,
+                        metrics='all', 
+                        test_name_list=test_name_list,
+                        render_all=cfg.params.RENDER_ALL,
+                        pck_thresh=cfg.params.PCK_THRESH, 
+                        return_results=True)
+    # prepare output mesh
+    mesh = all_results[0]['mesh_posed']
+    mesh.apply_transform([[-1, 0, 0, 0],
+                            [0, -1, 0, 0],
+                            [0, 0, 1, 1],
+                            [0, 0, 0, 1]])
+    result_path = os.path.join(save_imgs_path, test_name_list[0] + '_z')
+    mesh.export(file_obj=result_path + '.glb')
+    result_gltf = result_path + '.glb'
+    return result_gltf
+# -------------------------------------------------------------------------------------------------------------------- #
+
+
+def run_complete_inference(img_path_or_img, crop_choice):
+    # depending on crop_choice: run faster r-cnn or take the input image directly
+    if crop_choice == "input image is cropped":
+        if isinstance(img_path_or_img, str):
+            img = cv2.imread(img_path_or_img)
+            output_interm_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        else:
+            output_interm_image = img_path_or_img
+        output_interm_bbox = None
+    else:
+        output_interm_image, output_interm_bbox = run_bbox_inference(img_path_or_img.copy())
+    # run barc inference
+    result_gltf = run_barc_inference(img_path_or_img, output_interm_bbox)
+    # add white border to image for nicer alignment
+    output_interm_image_vis = np.concatenate((255*np.ones_like(output_interm_image), output_interm_image, 255*np.ones_like(output_interm_image)), axis=1)
+    return [result_gltf, result_gltf, output_interm_image_vis]
+
+
+
+
+########################################################################################################################
+
+# see: https://huggingface.co/spaces/radames/PIFu-Clothed-Human-Digitization/blob/main/PIFu/spaces.py
+
+description = '''
+# BARC
+
+#### Project Page
+* https://barc.is.tue.mpg.de/
+
+#### Description
+This is a demo for BARC (*B*reed *A*ugmented *R*egression using *C*lassification). 
+You can either submit a cropped image or choose the option to run a pretrained Faster R-CNN in order to obtain a bounding box.  
+Please have a look at the examples below.
+<details>
+
+<summary>More</summary>
+
+#### Citation
+
+```
+@inproceedings{BARC:2022,
+    title = {BARC}: Learning to Regress {3D} Dog Shape from Images by Exploiting Breed Information,
+    author = {Rueegg, Nadine and Zuffi, Silvia and Schindler, Konrad and Black, Michael J.},
+    booktitle = {Proceedings IEEE Conf. on Computer Vision and Pattern Recognition (CVPR)},
+    year = {2022}
+}
+```
+
+#### Image Sources (Examples)
+* Stanford extra image dataset
+* Images from google search engine
+    * https://www.dogtrainingnation.com/wp-content/uploads/2015/02/keep-dog-training-sessions-short.jpg
+    * https://thumbs.dreamstime.com/b/hund-und-seine-neue-hundeh%C3%BCtte-36757551.jpg
+    * https://www.mydearwhippet.com/wp-content/uploads/2021/04/whippet-temperament-2.jpg
+    * https://media.istockphoto.com/photos/ibizan-hound-at-the-shore-in-winter-picture-id1092705644?k=20&m=1092705644&s=612x612&w=0&h=ppwg92s9jI8GWnk22SOR_DWWNP8b2IUmLXSQmVey5Ss=
+
+
+</details>
+'''
+
+
+
+
+
+
+example_images = sorted(glob.glob(os.path.join(os.path.dirname(__file__), '../', 'datasets', 'test_image_crops', '*.jpg')) + glob.glob(os.path.join(os.path.dirname(__file__), '../', 'datasets', 'test_image_crops', '*.png')))  
+random.shuffle(example_images)
+examples = []
+for img in example_images:
+    if os.path.basename(img)[:2] == 'z_':
+        examples.append([img, "use Faster R-CNN to get a bounding box"])
+    else:
+        examples.append([img, "input image is cropped"])
+
+demo = gr.Interface(
+    fn=run_complete_inference,
+    description=description,
+    # inputs=gr.Image(type="filepath", label="Input Image"),
+    inputs=[gr.Image(label="Input Image"),
+        gr.Radio(["input image is cropped", "use Faster R-CNN to get a bounding box"], value="use Faster R-CNN to get a bounding box", label="Crop Choice"),
+    ],
+    outputs=[
+        gr.Model3D(
+            clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Model"),
+        gr.File(label="Download 3D Model"),
+        gr.Image(label="Bounding Box (Faster R-CNN prediction)"),
+
+    ],
+    examples=examples,
+    thumbnail="barc_thumbnail.png",
+    allow_flagging="never",
+    cache_examples=False,       # True
+    examples_per_page=14,
+)
+
+demo.launch()