Martin Tomov commited on
Commit
c0ddd70
1 Parent(s): 253fe94

color channels

Browse files
Files changed (1) hide show
  1. app.py +3 -8
app.py CHANGED
@@ -131,18 +131,13 @@ def mask_to_min_max(mask: np.ndarray) -> Tuple[int, int, int, int]:
131
  return x.min(), y.min(), x.max(), y.max()
132
 
133
  def extract_and_paste_insect(original_image: np.ndarray, detection: DetectionResult, background: np.ndarray) -> None:
134
- mask = detection.mask
135
  xmin, ymin, xmax, ymax = mask_to_min_max(mask)
136
  insect_crop = original_image[ymin:ymax, xmin:xmax]
137
  mask_crop = mask[ymin:ymax, xmin:xmax]
138
 
139
- # Ensure the insect retains its original color by directly copying the pixel values
140
- insect = np.zeros_like(insect_crop)
141
- for c in range(3):
142
- insect[:, :, c] = insect_crop[:, :, c] * (mask_crop / 255)
143
-
144
- for c in range(3):
145
- background[ymin:ymax, xmin:xmax, c] = background[ymin:ymax, xmin:xmax, c] * (1 - mask_crop / 255) + insect[:, :, c]
146
 
147
  def create_yellow_background_with_insects(image: np.ndarray, detections: List[DetectionResult]) -> np.ndarray:
148
  yellow_background = np.full((image.shape[0], image.shape[1], 3), (0, 255, 255), dtype=np.uint8) # BGR for yellow
 
131
  return x.min(), y.min(), x.max(), y.max()
132
 
133
  def extract_and_paste_insect(original_image: np.ndarray, detection: DetectionResult, background: np.ndarray) -> None:
134
+ mask = detection.mask.astype(bool)
135
  xmin, ymin, xmax, ymax = mask_to_min_max(mask)
136
  insect_crop = original_image[ymin:ymax, xmin:xmax]
137
  mask_crop = mask[ymin:ymax, xmin:xmax]
138
 
139
+ for c in range(3): # Loop over color channels
140
+ background[ymin:ymax, xmin:xmax, c][mask_crop] = insect_crop[:, :, c][mask_crop]
 
 
 
 
 
141
 
142
  def create_yellow_background_with_insects(image: np.ndarray, detections: List[DetectionResult]) -> np.ndarray:
143
  yellow_background = np.full((image.shape[0], image.shape[1], 3), (0, 255, 255), dtype=np.uint8) # BGR for yellow