NDVI Drone with SAM2 segmentation

demo_tools.py

import json
import colorsys
import cv2
import numpy as np

def sam_results(file):
    with open(file, 'r') as file:
        lines = file.readlines()
        data = []
        for line in lines:
            if line.startswith("{"):
                line = line.rstrip("--frame\n")
                try:
                    data.append(json.loads(line))
                except json.JSONDecodeError:
                    print("Error decoding JSON from line:", line)
        return data

def generate_colors(num_colors, saturation=0.9, lightness=0.5):
    colors = []
    hue = 1.0 / num_colors
    for i in range(num_colors):
        rgb = colorsys.hls_to_rgb(hue * i, lightness, saturation)
        colors.append(tuple(int(c * 255) for c in rgb))
    return colors

def get_contours(mask):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    lines = []
    for contour in contours:
        # Add the first point to the end to close the loop
        closed_contour = np.concatenate((contour, [contour[0]]), axis=0)
        lines.append(np.squeeze(closed_contour, axis=1))
    return lines

def calc_ndvi(color, nir):
    nir_float = cv2.cvtColor(nir, cv2.COLOR_BGR2GRAY).astype(float)
    _, _, band_red = cv2.split(color)
    band_red = band_red.astype(float)

    diff = nir_float - band_red
    sum = nir_float + band_red
    np.seterr(divide='ignore', invalid='ignore')
    return np.divide(diff, sum)

main.py

import cv2
import numpy as np
import rerun as rr
from pycocotools import mask as maskUtils
from demo_tools import sam_results, generate_colors, get_contours, calc_ndvi

caps = {}
for name in ['color', 'ir', 'ndvi_colorized']:
    caps[name] = cv2.VideoCapture(f'./{name}.mp4')

# Generate random colors
colors = np.concatenate((generate_colors(6, saturation=0.7, lightness=0.5), np.full((6, 1), 50)), axis=1)

# Run & initialize ReRun viewer
rr.init('NDVI /w SAM2', spawn=True)

# Prepare annotation context (for segmentation)
annotationContext = [(0, "Background", (0, 0, 0, 0))]
for i, color in enumerate(colors):
    annotationContext.append((i + 1, f"Field {i + 1}", color))
    rr.log(f"NDVI_Average/Field{i+1}", rr.SeriesLine(color=color, name=f"Field {i+1}", width=3), static=True)
rr.log("Color", rr.AnnotationContext(annotationContext), timeless=True)

sam_data = [sam_results(f"results{i}.txt") for i in range(1, 3)]
def get_sam_output(i):
    def match_i(arr):
        for obj in arr:
            if obj['frame_index'] == i:
                return obj
    return [match_i(arr) for arr in sam_data]

t = 0
size = (800, 1280)
for frame_idx in range(len(sam_data[0])):
    rr.set_time_sequence("step", t)
    t += 1

    frames = {} # Read all frames
    for name, cap in caps.items():
        ret, frame = cap.read()
        if not ret: break
        frames[name] = frame

    ndvi = calc_ndvi(frames['color'], frames['ir'])

    segmentations = np.zeros(size)
    for i, data in enumerate(get_sam_output(frame_idx)):
        for result in data.get("results", []):
            field_num = result['object_id']+i*3 # 3 segmentations per file
            # Decode the RLE mask
            mask = np.array(maskUtils.decode(result["mask"]), dtype=np.uint8)
            # Set full_mask to num where mask
            segmentations[mask == 1] = field_num + 1 # 0 == Background

            line_strips = get_contours(mask)
            rr.log(f"Color/Contours{field_num + 1}", rr.LineStrips2D(line_strips, colors=colors[field_num], labels=f"Field {field_num + 1}"))
            rr.log(f"NDVI/Color/Contours{field_num + 1}", rr.LineStrips2D(line_strips, colors=colors[field_num], labels=f"Field {field_num + 1}"))

            mean_ndvi = np.mean(ndvi[mask == 1])
            rr.log(f"NDVI_Average/Field{field_num + 1}", rr.Scalar(mean_ndvi))

    rr.log("Color/Image", rr.Image(frames['color'][..., ::-1]))
    rr.log("NDVI/Color", rr.Image(frames['ndvi_colorized'][..., ::-1]))
    rr.log("Color/Mask", rr.SegmentationImage(segmentations))
import json
import colorsys
import cv2
import numpy as np

def sam_results(file):
    with open(file, 'r') as file:
        lines = file.readlines()
        data = []
        for line in lines:
            if line.startswith("{"):
                line = line.rstrip("--frame\n")
                try:
                    data.append(json.loads(line))
                except json.JSONDecodeError:
                    print("Error decoding JSON from line:", line)
        return data

def generate_colors(num_colors, saturation=0.9, lightness=0.5):
    colors = []
    hue = 1.0 / num_colors
    for i in range(num_colors):
        rgb = colorsys.hls_to_rgb(hue * i, lightness, saturation)
        colors.append(tuple(int(c * 255) for c in rgb))
    return colors

def get_contours(mask):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    lines = []
    for contour in contours:
        # Add the first point to the end to close the loop
        closed_contour = np.concatenate((contour, [contour[0]]), axis=0)
        lines.append(np.squeeze(closed_contour, axis=1))
    return lines

def calc_ndvi(color, nir):
    nir_float = cv2.cvtColor(nir, cv2.COLOR_BGR2GRAY).astype(float)
    _, _, band_red = cv2.split(color)
    band_red = band_red.astype(float)

    diff = nir_float - band_red
    sum = nir_float + band_red
    np.seterr(divide='ignore', invalid='ignore')
    return np.divide(diff, sum)