Gallai commited on
Commit
33fad17
·
verified ·
1 Parent(s): 5aadd0d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +71 -0
app.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import joblib
3
+ import cv2
4
+ import numpy as np
5
+ from PIL import Image
6
+ from tensorflow.keras.models import load_model
7
+
8
+ # Define paths to models and load the scaler
9
+ model_paths = {
10
+ "Regressor_decision_tree": "multioutput_regressor_decision_tree.joblib",
11
+ "Regressor_ridge": "regressor_ridge.joblib",
12
+ "Regressor_elastic_net": "elastic_net_model.joblib",
13
+ "NN_6_Layers": "NN_Layers_6.keras",
14
+ "CNN": "cnn_model_bigger.keras",
15
+ "CNN_with_reductions": "cnn_model_bigger_with_reductions.keras"
16
+ }
17
+ scaler = joblib.load("scaler.joblib")
18
+
19
+ # Function to load models based on file extension
20
+ def load_model_by_type(path):
21
+ if path.endswith('.joblib'):
22
+ return joblib.load(path)
23
+ elif path.endswith('.keras'):
24
+ return load_model(path)
25
+ else:
26
+ raise ValueError(f"Unsupported file extension for file {path}")
27
+
28
+ # Load models with appropriate method
29
+ models = {name: load_model_by_type(path) for name, path in model_paths.items()}
30
+
31
+
32
+ def detect_objects(image, model_name):
33
+ model = models[model_name]
34
+
35
+ # Assuming Gradio passes image as a numpy array and checking if conversion is needed
36
+ if image.ndim == 3 and image.shape[2] == 3: # If the image is RGB
37
+ image_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) # Convert to grayscale using OpenCV
38
+ else:
39
+ image_gray = image # Use the image as is if already grayscale
40
+
41
+ # Check if the model requires CNN specific preprocessing
42
+ if model_name in ["CNN_model", "CNN_with_reductions"]:
43
+ image_processed = np.array(image_gray)
44
+ image_processed = image_processed.reshape(1, image_gray.shape[0], image_gray.shape[1], 1)
45
+ image_processed = image_processed.astype('float32')
46
+ image_processed /= 255 # Normalize pixel values
47
+ else:
48
+ # Assuming other models might expect flattened, scaled input
49
+ image_processed = image_gray.flatten().reshape(1, -1)
50
+ image_processed = scaler.transform(image_processed)
51
+
52
+ # Make prediction
53
+ predictions = model.predict(image_processed)
54
+ x, y, width, height = predictions[0]
55
+
56
+ # Draw bounding box on a copy of the original image (converted back to RGB for color drawing)
57
+ original_image_rgb = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) # Ensure image is in RGB
58
+ cv2.rectangle(original_image_rgb, (int(x), int(y)), (int(x + width), int(y + height)), (0, 255, 0), 2)
59
+
60
+ return Image.fromarray(original_image_rgb)
61
+
62
+ # Gradio interface setup
63
+ iface = gr.Interface(
64
+ fn=detect_objects,
65
+ inputs=[gr.components.Image(), gr.components.Dropdown(list(model_paths.keys()))],
66
+ outputs=gr.components.Image(),
67
+ title="Object Detection",
68
+ description="Select a model and upload an image to detect objects."
69
+ )
70
+
71
+ iface.launch(show_error=True, share=True, debug=True)