hololens
commited on
Commit
β’
ec41162
1
Parent(s):
326f023
init commit
Browse files
app.py
ADDED
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import logging
|
3 |
+
import gradio as gr
|
4 |
+
from PIL import Image as PILImg
|
5 |
+
from iteach_toolkit.DHYOLO import DHYOLODetector
|
6 |
+
|
7 |
+
# Configure logging
|
8 |
+
logging.basicConfig(level=logging.INFO)
|
9 |
+
logger = logging.getLogger(__name__)
|
10 |
+
|
11 |
+
def detect_objects(selected_model, input_image, conf, imu_threshold, detections):
|
12 |
+
try:
|
13 |
+
model_path = model_options[selected_model]
|
14 |
+
dhyolo = DHYOLODetector(model_path)
|
15 |
+
|
16 |
+
# Save the input image temporarily to a path for processing
|
17 |
+
input_image_path = "dhyolo_temp_input_image.jpg"
|
18 |
+
input_image.save(input_image_path)
|
19 |
+
|
20 |
+
# Perform prediction on the image
|
21 |
+
orig_image, detections = dhyolo.predict(input_image_path, conf, imu_threshold, detections)
|
22 |
+
|
23 |
+
# Log the detections
|
24 |
+
logger.info("Detections: %s", detections)
|
25 |
+
|
26 |
+
# Plot the bounding boxes on the original image
|
27 |
+
orig_image, image_with_bboxes = dhyolo.plot_bboxes(attach_watermark=True)
|
28 |
+
|
29 |
+
# Convert the image (with bounding boxes) from a NumPy array to a PIL Image for display.
|
30 |
+
pil_img_with_bboxes = PILImg.fromarray(image_with_bboxes)
|
31 |
+
|
32 |
+
return input_image, pil_img_with_bboxes
|
33 |
+
|
34 |
+
except FileNotFoundError as e:
|
35 |
+
logger.error("File not found: %s", e)
|
36 |
+
return None, None
|
37 |
+
except Exception as e:
|
38 |
+
logger.error("An error occurred: %s", e)
|
39 |
+
return None, None
|
40 |
+
|
41 |
+
def load_test_images():
|
42 |
+
"""Load images from the test_imgs directory."""
|
43 |
+
test_imgs_dir = os.path.join(os.getcwd(), "test_imgs")
|
44 |
+
logger.info("Loading images from: %s", test_imgs_dir) # Log the directory path
|
45 |
+
return [f for f in os.listdir(test_imgs_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
|
46 |
+
|
47 |
+
# Gradio interface
|
48 |
+
def create_interface():
|
49 |
+
with gr.Blocks() as demo:
|
50 |
+
# Center the title using HTML
|
51 |
+
gr.Markdown("<h1 style='text-align: center;'>πͺπ DHYOLO DoorHandle Object Detection</h1>")
|
52 |
+
|
53 |
+
# Add paper description with correct hyperlink syntax
|
54 |
+
gr.Markdown("<h2 style='text-align: center;'>π iTeach: Interactive Teaching for Robot Perception using Mixed Reality</h2>")
|
55 |
+
|
56 |
+
# Add the project link using HTML <a> tag
|
57 |
+
gr.Markdown("<h2 style='text-align: center;'>π Project Link: <a href='https://irvlutd.github.io/iTeach/' target='_blank'>iTeach Project</a></h2>")
|
58 |
+
|
59 |
+
with gr.Row():
|
60 |
+
with gr.Column():
|
61 |
+
# Load a default image for input
|
62 |
+
default_image_path = os.path.join(os.getcwd(), "test_imgs", "irvl-ters.jpg") # Update with the correct default image name
|
63 |
+
|
64 |
+
# Attempt to open the default image and log if it fails
|
65 |
+
try:
|
66 |
+
default_image = PILImg.open(default_image_path)
|
67 |
+
except FileNotFoundError:
|
68 |
+
logger.error("Default image not found at: %s", default_image_path)
|
69 |
+
default_image = None # Set to None or a placeholder image
|
70 |
+
|
71 |
+
input_image = gr.Image(type="pil", label="Input Image", value=default_image)
|
72 |
+
|
73 |
+
# Pretrained model paths
|
74 |
+
cwd = os.getcwd()
|
75 |
+
global model_options
|
76 |
+
model_options = {
|
77 |
+
"dh-yolo-v1-pb-ddf-524": f'{cwd}/pretrained_ckpts/dh-yolo-v1-pb-ddf-524.pt',
|
78 |
+
"dh-yolo-exp27-pb-1008": f'{cwd}/pretrained_ckpts/dh-yolo-exp27-pb-1008.pt',
|
79 |
+
"dh-yolo-exp31-pb-1532": f'{cwd}/pretrained_ckpts/dh-yolo-exp-31-pb-1532.pt',
|
80 |
+
"dh-yolo-exp31-pl-1532": f'{cwd}/pretrained_ckpts/dh-yolo-exp-31-pl-1532.pt'
|
81 |
+
}
|
82 |
+
|
83 |
+
model_path = gr.Dropdown(choices=list(model_options.keys()), label="Select Pretrained Model", value="dh-yolo-v1-pb-ddf-524")
|
84 |
+
|
85 |
+
conf = gr.Slider(label="Confidence Threshold", minimum=0.0, maximum=1.0, step=0.01, value=0.5)
|
86 |
+
imu_threshold = gr.Slider(label="IMU Threshold", minimum=0.0, maximum=1.0, step=0.01, value=0.5)
|
87 |
+
detections = gr.Slider(label="Max number of Detections", minimum=1, maximum=100, step=1, value=10)
|
88 |
+
|
89 |
+
with gr.Column():
|
90 |
+
output_image = gr.Image(label="Output Image with DH-YOLO Detections", type="pil")
|
91 |
+
|
92 |
+
detect_button = gr.Button("Run")
|
93 |
+
|
94 |
+
# Detect button functionality
|
95 |
+
detect_button.click(detect_objects, inputs=[model_path, input_image, conf, imu_threshold, detections],
|
96 |
+
outputs=[input_image, output_image])
|
97 |
+
|
98 |
+
return demo
|
99 |
+
|
100 |
+
if __name__ == "__main__":
|
101 |
+
demo = create_interface()
|
102 |
+
demo.launch()
|
pretrained_ckpts/dh-yolo-exp-31-pb-1532.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:625a45abe7ae0e7d52a7528915f45408a93865fc9711e42663a37ce61890e68f
|
3 |
+
size 42181652
|
pretrained_ckpts/dh-yolo-exp-31-pl-1532.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:a94b4f8d0f6d965a9de0cf3e22ce3f478bd198eca858e5cce7522393ca07190f
|
3 |
+
size 42181652
|
pretrained_ckpts/dh-yolo-exp27-pb-1008.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:a419efeb648285e265195fab71d1e23bc8a95b386f1b10e5bc80abedb26a1859
|
3 |
+
size 42181652
|
pretrained_ckpts/dh-yolo-v1-pb-ddf-524.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:b0179ed926a58f2c379dda9fbfa15cbca23cefc058554efd9325e4d5f5531e81
|
3 |
+
size 42181652
|
requirements.txt
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
gradio
|
2 |
+
Pillow
|
3 |
+
iteach_toolkit
|
4 |
+
opencv-python
|