consolidating codes

app.py

@@ -1,72 +1,84 @@
 # Copied megadetector section from https://huggingface.co/spaces/hlydecker/MegaDetector_v5 
+# Copied from https://huggingface.co/spaces/Neslihan/megadetector_dlcmodels/blob/main/app.py 
 
-# %%
-#all imports
+print("before import")
 import gradio as gr
-import torch
-import torchvision
+import json
+import os
+import matplotlib.pyplot as plt
 import numpy as np
 from PIL import Image
-import matplotlib.pyplot as plt
+from dlclive import DLCLive, Processor
+from numpy import savetxt
+import PIL
+print("after import")
+
+# A method that allows using dlc live, but with different models, saves poses, and plots the poses onto image
+def dlclive_pose(model, crop_np, crop, index,dlc_proc):
+    dlc_live = DLCLive(model, processor=dlc_proc) 
+    dlc_live.init_inference(crop_np)
+    keypts = dlc_live.get_pose(crop_np) 
+    xpose = []
+    ypose = []
+    for key in keypts[:,2]:
+       # if key > 0.05: # which value do we need here?
+            i = np.where(keypts[:,2]==key)
+            xpose.append(keypts[i,0])
+            ypose.append(keypts[i,1])
+    plt.imshow(crop)
+    plt.scatter(xpose[:], ypose[:], 40, color='cyan')
+
+    canvas = plt.gca().figure.canvas
+    canvas.draw()
+    image = PIL.Image.frombytes('RGB', canvas.get_width_height(), canvas.tostring_rgb())
+
+    plt.clf()
+    return image
+
+def classify_image(img, file):
+
+    primate_face_model = 'model_weights/DLC_FacialLandmarks_resnet_50_iteration-1_shuffle-1'
+    human_model = 'model_weights/DLC_human_dancing_resnet_101_iteration-0_shuffle-1' 
+
+    with open(file.name, 'r') as f:
+        detection_results = json.load(f)
 
-# %%
-# Loads a model from github repo, but you need to have the model 
-model = torch.hub.load('ultralytics/yolov5', 'custom', "https://huggingface.co/spaces/vchiang001/MegaDetector_DLC_pose/blob/main/md_v5b.0.0.pt", force_reload=True)
-
-# %% 
-
-#not sure if we need to resize... maybe just the origin image and see? 
-""" def yolo(im): #size=640):
-g = (size / max(im.size))  # gain
-    im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS)  # resize 
-    
-    model = torch.hub.load('ultralytics/yolov5', 'custom', "https://huggingface.co/spaces/vchiang001/MegaDetector_DLC_pose/blob/main/md_v5b.0.0.pt", force_reload=True)
-    
-    results = model(im)  # inference
-    results.render()  # updates results.imgs with boxes and labels
-    return Image.fromarray(results.imgs[0])  """
-
-def yolo(im): #size=640):
-    model = torch.hub.load('ultralytics/yolov5', 'custom', "https://huggingface.co/spaces/vchiang001/MegaDetector_DLC_pose/blob/main/md_v5b.0.0.pt", force_reload=True)
-    results = model(im)  # inference
-    results.render()  # updates results.imgs with boxes and labels
-    return Image.fromarray(results.imgs[0])
-
-# %% dlc-live function
-def dlc(im):
-    from dlclive import DLCLive, Processor
-    imarray = np.asarray(im)
-    dlc_proc = Processor()
-    dlc_live = DLCLive("model/", processor=dlc_proc)
-    dlc_live.init_inference(imarray)
-    keypts = dlc_live.get_pose(imarray)
-    return Image.fromarray(keypts)
-"""
-    from dlclive import DLCLive, Processor
-    imarray = np.asarray(im)
     dlc_proc = Processor()
-    dlc_live = DLCLive("model/", processor=dlc_proc)
-    dlc_live.init_inference(imarray)
-    keypts = dlc_live.get_pose(imarray)
-    plt.imshow(im)
-    plt.scatter(keypts[:,0], keypts[:,1], 40, color='cyan')
-    #return Image.fromarray(results.imgs[0])
-    """
-    #return plt.show()
-    #return plt.imshow(im)
-
-# %% 
-#inputs = [image, chosen_model, size]
-#this is where you show the image as interface 
-inputs = gr.inputs.Image(type="pil", label="Input Image")
-outputs = gr.outputs.Image(type="pil", label="Output Image")
-
-
-# %% 
-title = "MegaDetector with DeepLabCut pose estimation"
-description = "Detect and identify animals, people and vehicles in camera trap images followed by generating poses for humans and animals"
-article = "<p style='text-align: center'>Detect and identify animals, people and vehicles in camera trap images followed by generating poses for humans and animals</a></p>"
-
-# %% 
-#running the actual 
-gr.Interface(dlc, inputs, outputs, title=title, description=description, article=article, theme="huggingface").launch(enable_queue=True)
+
+    # Assuming there is only 1 detection on the output
+    img_data = detection_results["images"][0]
+
+    output_images = []
+
+    for detections_dict in img_data["detections"]:
+        index = img_data["detections"].index(detections_dict)
+        if detections_dict["conf"] > 0.8: 
+            x1, y1,w_box, h_box = detections_dict["bbox"]
+            ymin,xmin,ymax, xmax = y1, x1, y1 + h_box, x1 + w_box
+            
+            imageWidth=img.size[0]
+            imageHeight= img.size[1]
+            area = (xmin * imageWidth, ymin * imageHeight, xmax * imageWidth,
+                    ymax * imageHeight)
+            crop = img.crop(area)
+            crop_np = np.asarray(crop)
+            
+            if detections_dict["category"] == "1":
+                selected_model = primate_face_model
+            elif detections_dict["category"] == "2":
+                selected_model = human_model
+
+            # Until we know how to dynamically add output element to gradio, just return the first image
+            output_images.append(dlclive_pose(selected_model, crop_np, crop, index, dlc_proc))
+
+    return output_images[0], output_images[1] # lol
+
+input_image = gr.inputs.Image(type="pil", label="Input Image")
+input_file = gr.inputs.File(label="output.json")
+
+# Fake it till we make it, we know our example has 2 outputs
+outputs = [gr.outputs.Image(type="pil", label="Output Image"), gr.outputs.Image(type="pil", label="Output Image")]
+
+gr.Interface(fn=classify_image, inputs=[input_image, input_file], outputs=outputs, theme="huggingface").launch()
+
+