Spaces:
Sleeping
Sleeping
Update model_1.py
Browse files- model_1.py +6 -4
model_1.py
CHANGED
@@ -48,10 +48,10 @@ color_mapping = {
|
|
48 |
'Empty class': (0, 255, 0) # Green for empty
|
49 |
}
|
50 |
|
51 |
-
def perform_detection(image):
|
52 |
imH, imW, _ = image.shape
|
53 |
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
54 |
-
image_resized = cv2.resize(image_rgb,
|
55 |
input_data = np.expand_dims(image_resized, axis=0)
|
56 |
|
57 |
if floating_model:
|
@@ -86,8 +86,10 @@ def perform_detection(image):
|
|
86 |
return image
|
87 |
|
88 |
def detect_image(input_image):
|
|
|
|
|
89 |
image = np.array(input_image)
|
90 |
-
result_image = perform_detection(image)
|
91 |
return Image.fromarray(result_image)
|
92 |
|
93 |
def detect_video(input_video):
|
@@ -99,7 +101,7 @@ def detect_video(input_video):
|
|
99 |
if not ret:
|
100 |
break
|
101 |
|
102 |
-
result_frame = perform_detection(frame)
|
103 |
frames.append(result_frame)
|
104 |
|
105 |
cap.release()
|
|
|
48 |
'Empty class': (0, 255, 0) # Green for empty
|
49 |
}
|
50 |
|
51 |
+
def perform_detection(image, target_size=(640, 640)):
|
52 |
imH, imW, _ = image.shape
|
53 |
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
54 |
+
image_resized = cv2.resize(image_rgb, target_size)
|
55 |
input_data = np.expand_dims(image_resized, axis=0)
|
56 |
|
57 |
if floating_model:
|
|
|
86 |
return image
|
87 |
|
88 |
def detect_image(input_image):
|
89 |
+
# Resize the image to 640x640 for faster processing
|
90 |
+
target_size = (640, 640)
|
91 |
image = np.array(input_image)
|
92 |
+
result_image = perform_detection(image, target_size)
|
93 |
return Image.fromarray(result_image)
|
94 |
|
95 |
def detect_video(input_video):
|
|
|
101 |
if not ret:
|
102 |
break
|
103 |
|
104 |
+
result_frame = perform_detection(frame, (640, 640))
|
105 |
frames.append(result_frame)
|
106 |
|
107 |
cap.release()
|