sabaridsnfuji commited on
Commit
f144fcd
1 Parent(s): e00ea77

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +109 -129
app.py CHANGED
@@ -1,129 +1,109 @@
1
- # -*- coding: utf-8 -*-
2
- """
3
- Created on Wed Nov 13 18:37:31 2024
4
- @author: sabar
5
- """
6
-
7
- import gradio as gr
8
- import cv2
9
- import numpy as np
10
- import os
11
- import json
12
- from openvino.runtime import Core # Assuming you're using OpenVINO
13
- from tf_post_processing import non_max_suppression # Assuming this is defined elsewhere
14
- from PIL import Image
15
-
16
- # Load the OpenVINO model
17
- classification_model_xml = "./model/best.xml"
18
- core = Core()
19
- config = {
20
- "INFERENCE_NUM_THREADS": 2,
21
- "ENABLE_CPU_PINNING": True
22
- }
23
- model = core.read_model(model=classification_model_xml)
24
- compiled_model = core.compile_model(model=model, device_name="CPU", config=config)
25
-
26
- label_to_class_text = {
27
- 0: 'range',
28
- 1: 'entry door',
29
- 2: 'kitchen sink',
30
- 3: 'bathroom sink',
31
- 4: 'toilet',
32
- 5: 'double folding door',
33
- 6: 'window',
34
- 7: 'shower',
35
- 8: 'bathtub',
36
- 9: 'single folding door',
37
- 10: 'dishwasher',
38
- 11: 'refrigerator'
39
- }
40
-
41
- # Function to perform inference
42
- def predict_image(image):
43
- # Convert the Pillow image to a NumPy array and BGR format for OpenCV
44
- image = np.array(image.convert("RGB"))
45
- image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
46
-
47
- # Resize, preprocess, and reshape the input image
48
- img_size = 960
49
- resized_image = cv2.resize(image, (img_size, img_size)) / 255.0
50
- resized_image = resized_image.transpose(2, 0, 1)
51
- reshaped_image = np.expand_dims(resized_image, axis=0).astype(np.float32)
52
-
53
- im_height, im_width, _ = image.shape
54
- output_numpy = compiled_model(reshaped_image)[0]
55
- results = non_max_suppression(output_numpy, conf_thres=0.2, iou_thres=0.6, max_wh=15000)[0]
56
-
57
- # Prepare output paths
58
- output_path = "./output_file_train/"
59
- output_image_folder = os.path.join(output_path, "images_alienware_openvino/")
60
- os.makedirs(output_image_folder, exist_ok=True)
61
-
62
- output_json_folder = os.path.join(output_path, "json_output/")
63
- os.makedirs(output_json_folder, exist_ok=True)
64
-
65
- predictions = []
66
-
67
- # Draw boxes and collect prediction data
68
- for result in results:
69
- boxes = result[:4]
70
- prob = result[4]
71
- classes = int(result[5])
72
-
73
- x1, y1, x2, y2 = np.uint16([
74
- boxes[0] * im_width,
75
- boxes[1] * im_height,
76
- boxes[2] * im_width,
77
- boxes[3] * im_height
78
- ])
79
-
80
- if prob > 0.2:
81
- cv2.rectangle(image, (x1, y1), (x2, y2), (255, 255, 0), 2)
82
- label_text = f"{classes} {round(prob, 2)}"
83
- cv2.putText(image, label_text, (x1, y1), 0, 0.5, (0, 255, 0), 2)
84
-
85
- # Store prediction info in a JSON-compatible format
86
- predictions.append({
87
- "class": label_to_class_text[classes],
88
- "probability": round(float(prob), 2),
89
- "coordinates": {
90
- "xmin": int(x1),
91
- "ymin": int(y1),
92
- "xmax": int(x2),
93
- "ymax": int(y2)
94
- }
95
- })
96
-
97
- # Save the processed image
98
- output_image_path = os.path.join(output_image_folder, "result_image.jpg")
99
- cv2.imwrite(output_image_path, image)
100
-
101
- # Convert predictions to a formatted string
102
- predictions_str = json.dumps(predictions, indent=4)
103
-
104
- return image, predictions_str
105
-
106
- # Define sample images for user convenience
107
- sample_images = [
108
- "./sample/10_2.jpg",
109
- "./sample/10_10.jpg",
110
- "./sample/10_12.jpg"
111
- ]
112
-
113
- # Set up Gradio interface
114
- def gradio_interface(image):
115
- output_image, predictions_str = predict_image(image)
116
- return output_image, predictions_str
117
-
118
- # Create the Gradio Interface
119
- gr_interface = gr.Interface(
120
- fn=gradio_interface,
121
- inputs=gr.Image(label="Upload or Select an Image", type="pil", examples=sample_images),
122
- outputs=[gr.Image(label="Result Image"), gr.Textbox(label="Predictions JSON")],
123
- title="House CAD Design Object Detection",
124
- description="Upload a CAD design image of a house to detect objects with bounding boxes and probabilities."
125
- )
126
-
127
- # Launch the Gradio interface if run as main
128
- if __name__ == "__main__":
129
- gr_interface.launch()
 
1
+ import gradio as gr
2
+ import cv2
3
+ import numpy as np
4
+ import os
5
+ import json
6
+ from openvino.runtime import Core # Assuming you're using OpenVINO
7
+ from tqdm import tqdm
8
+ from PIL import Image
9
+
10
+ from tf_post_processing import non_max_suppression #,optimized_object_detection
11
+
12
+ # Load the OpenVINO model
13
+ classification_model_xml = r"D:\Research\Auto_CAD_detection\runs\detect\yolo11_Auto_CAD_detection_params6\weights\best_openvino_model\best.xml"
14
+ core = Core()
15
+ config = {
16
+ "INFERENCE_NUM_THREADS": 2,
17
+ "ENABLE_CPU_PINNING": True
18
+ }
19
+ model = core.read_model(model=classification_model_xml)
20
+ compiled_model = core.compile_model(model=model, device_name="CPU", config=config)
21
+
22
+ label_to_class_text = {0: 'range',
23
+ 1: ' entry door',
24
+ 2: 'kitchen sink',
25
+ 3: ' bathroom sink',
26
+ 4: 'toilet',
27
+ 5: 'double folding door',
28
+ 6: 'window',
29
+ 7: 'shower',
30
+ 8: 'bathtub',
31
+ 9: 'single folding door',
32
+ 10: 'dishwasher',
33
+ 11: 'refrigerator'}
34
+
35
+ # Function to perform inference
36
+ def predict_image(image):
37
+ # Convert PIL Image to numpy array (OpenCV uses numpy arrays)
38
+ image = np.array(image)
39
+
40
+ # Resize, preprocess, and reshape the input image
41
+ img_size = 960
42
+ resized_image = cv2.resize(image, (img_size, img_size)) / 255.0
43
+ resized_image = resized_image.transpose(2, 0, 1)
44
+ reshaped_image = np.expand_dims(resized_image, axis=0).astype(np.float32)
45
+
46
+ im_height, im_width, _ = image.shape
47
+ output_numpy = compiled_model(reshaped_image)[0]
48
+ results = non_max_suppression(output_numpy, conf_thres=0.2, iou_thres=0.6, max_wh=15000)[0]
49
+
50
+ # Prepare output paths
51
+ predictions = []
52
+
53
+ # Draw boxes and collect prediction data
54
+ for result in results:
55
+ boxes = result[:4]
56
+ prob = result[4]
57
+ classes = int(result[5])
58
+
59
+ x1, y1, x2, y2 = np.uint16([
60
+ boxes[0] * im_width,
61
+ boxes[1] * im_height,
62
+ boxes[2] * im_width,
63
+ boxes[3] * im_height
64
+ ])
65
+
66
+ if prob > 0.2:
67
+ cv2.rectangle(image, (x1, y1), (x2, y2), (255, 255, 0), 2)
68
+ label_text = f"{classes} {round(prob, 2)}"
69
+ cv2.putText(image, label_text, (x1, y1), 0, 0.5, (0, 255, 0), 2)
70
+
71
+ # Store prediction info in a JSON-compatible format
72
+ predictions.append({
73
+ "class": label_to_class_text[classes],
74
+ "probability": round(float(prob), 2),
75
+ "coordinates": {
76
+ "xmin": int(x1),
77
+ "ymin": int(y1),
78
+ "xmax": int(x2),
79
+ "ymax": int(y2)
80
+ }
81
+ })
82
+
83
+ # Convert the processed image back to PIL Image for Gradio
84
+ pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
85
+
86
+ return pil_image, json.dumps(predictions, indent=4)
87
+
88
+ # Sample images for Gradio examples
89
+ # Define sample images for user convenience
90
+ sample_images = [
91
+ "./sample/10_2.jpg",
92
+ "./sample/10_10.jpg",
93
+ "./sample/10_12.jpg"
94
+ ]
95
+
96
+
97
+ # Gradio UI setup with examples
98
+ gr_interface = gr.Interface(
99
+ fn=predict_image,
100
+ inputs=gr.Image(type="pil"), # Updated to gr.Image for image input
101
+ outputs=[gr.Image(type="pil"), gr.Textbox()], # Updated to gr.Image and gr.Textbox
102
+ title="House CAD Design Object Detection",
103
+ description="Upload a CAD design image of a house to detect objects with bounding boxes and probabilities.",
104
+ examples=sample_images # Add the examples here
105
+ )
106
+
107
+ # Launch the Gradio interface if run as main
108
+ if __name__ == "__main__":
109
+ gr_interface.launch()