72 lines
2.0 KiB
Python
72 lines
2.0 KiB
Python
import cv2
|
|
import numpy as np
|
|
import os
|
|
|
|
def process_hay():
|
|
# Paths
|
|
input_path = 'assets/references/proizvodnja_sena.png'
|
|
output_dir = 'assets/DEMO_FAZA1/Items/Hay'
|
|
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
|
|
# Load image
|
|
img = cv2.imread(input_path)
|
|
if img is None:
|
|
print(f"Error: Could not load {input_path}")
|
|
return
|
|
|
|
# Convert to RGBA
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
|
|
|
|
# Define Green Key (assuming bright green background)
|
|
# Range for Chroma Key Green in HSV is best, but BGR is fine for simple flat green
|
|
# Pure Green in BGR is (0, 255, 0)
|
|
|
|
# Using HSV for better masking
|
|
hsv = cv2.cvtColor(img[:,:,:3], cv2.COLOR_BGR2HSV)
|
|
|
|
# Mask for Green (Hue 35-85 approx for standard green screen)
|
|
lower_green = np.array([35, 50, 50])
|
|
upper_green = np.array([85, 255, 255])
|
|
|
|
mask = cv2.inRange(hsv, lower_green, upper_green)
|
|
|
|
# Invert mask (we want non-green parts)
|
|
mask_inv = cv2.bitwise_not(mask)
|
|
|
|
# Set Alpha channel based on mask
|
|
# Everything green becomes transparent
|
|
img[:, :, 3] = mask_inv
|
|
|
|
# Find contours of objects
|
|
# Use the mask_inv (where 255 is the object)
|
|
contours, _ = cv2.findContours(mask_inv, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
print(f"Found {len(contours)} objects.")
|
|
|
|
count = 0
|
|
for cnt in contours:
|
|
x, y, w, h = cv2.boundingRect(cnt)
|
|
|
|
# Filter noise
|
|
if w < 10 or h < 10:
|
|
continue
|
|
|
|
# Crop
|
|
crop = img[y:y+h, x:x+w]
|
|
|
|
# Determine if it's a small pile or big bale based on size
|
|
# This is a heuristic. We'll label them by size.
|
|
area = w * h
|
|
label = "hay_piece"
|
|
|
|
# Save
|
|
output_filename = f"{output_dir}/hay_drop_{count}.png"
|
|
cv2.imwrite(output_filename, crop)
|
|
print(f"Saved {output_filename} (Size: {w}x{h})")
|
|
count += 1
|
|
|
|
if __name__ == "__main__":
|
|
process_hay()
|