Transfinity_bot / geometry.py
Nurbot's picture
Upload 5 files
2638527 verified
import cv2
import numpy as np
def extract_candles(image_path):
img = cv2.imread(image_path)
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
green_lower = np.array([45, 80, 80])
green_upper = np.array([75, 255, 255])
red_lower1 = np.array([0, 100, 80])
red_upper1 = np.array([10, 255, 255])
red_lower2 = np.array([160, 100, 80])
red_upper2 = np.array([179, 255, 255])
mask_green = cv2.inRange(hsv, green_lower, green_upper)
mask_red1 = cv2.inRange(hsv, red_lower1, red_upper1)
mask_red2 = cv2.inRange(hsv, red_lower2, red_upper2)
mask_red = cv2.bitwise_or(mask_red1, mask_red2)
contours_g, _ = cv2.findContours(mask_green, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
contours_r, _ = cv2.findContours(mask_red, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
candles = []
for c in contours_g + contours_r:
x, y, w, h = cv2.boundingRect(c)
if h > 10 and w > 2:
body = h / w
price = 0.19200 + (600 - y) * 0.00001 # fake scale
candles.append({
'x': x, 'y': y, 'w': w, 'h': h,
'body': body,
'type': 'green' if c in contours_g else 'red',
'price': round(price, 5)
})
candles = sorted(candles, key=lambda c: c['x'])
return candles[-10:]