brxerq commited on
Commit
21f6d14
1 Parent(s): 8966657

Update model_1.py

Browse files
Files changed (1) hide show
  1. 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, (width, height))
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()