camera-3P1200-PD-

color.py

#!/usr/bin/env python3
import gi
gi.require_version('Aravis', '0.8')
from gi.repository import Aravis
import numpy as np
import cv2
import sys
import time
import json
import os

class InspectionCamera:
    def __init__(self, profile_file='camera_profile.json'):
        try:
            self.camera = Aravis.Camera.new(None)
        except Exception as e:
            print(f"Error: Cannot connect to camera: {e}")
            sys.exit(1)
            
        self.stream = self.camera.create_stream(None, None)
        self.width, self.height = self.camera.get_sensor_size()
        self.profile_file = profile_file
        
        print(f"Camera connected: {self.width}x{self.height}")
        
        self.camera.set_region(0, 0, self.width, self.height)
        self.camera.set_exposure_time(50000.0)
        
        try:
            self.camera.set_string("BalanceWhiteAuto", "Off")
            self.set_hw_white_balance(3.0, 0.5)
        except Exception as e:
            print(f"White balance control error: {e}")
        
        # Default software white balance
        self.sw_red = 1.0
        self.sw_green = 0.9
        self.sw_blue = 1.2
        
        # Load saved profile if exists
        self.load_profile()
        
        self.reference_image = None
    
    def save_profile(self):
        """Save current white balance settings to file"""
        profile = {
            'sw_red': self.sw_red,
            'sw_green': self.sw_green,
            'sw_blue': self.sw_blue,
            'timestamp': time.time()
        }
        try:
            with open(self.profile_file, 'w') as f:
                json.dump(profile, f, indent=2)
            print(f"✓ Profile saved to {self.profile_file}")
            print(f"  R:{self.sw_red:.2f} G:{self.sw_green:.2f} B:{self.sw_blue:.2f}")
            return True
        except Exception as e:
            print(f"Error saving profile: {e}")
            return False
    
    def load_profile(self):
        """Load white balance settings from file"""
        if os.path.exists(self.profile_file):
            try:
                with open(self.profile_file, 'r') as f:
                    profile = json.load(f)
                self.sw_red = profile.get('sw_red', 1.0)
                self.sw_green = profile.get('sw_green', 1.0)
                self.sw_blue = profile.get('sw_blue', 1.0)
                print(f"✓ Loaded profile from {self.profile_file}")
                print(f"  R:{self.sw_red:.2f} G:{self.sw_green:.2f} B:{self.sw_blue:.2f}")
                return True
            except Exception as e:
                print(f"Error loading profile: {e}")
        else:
            print(f"No saved profile found at {self.profile_file}")
        return False
    
    def load_reference_image(self, filename):
        """Load reference image for color matching"""
        try:
            self.reference_image = cv2.imread(filename)
            if self.reference_image is not None:
                print(f"Loaded reference image: {filename}")
                avg_b = np.mean(self.reference_image[:,:,0])
                avg_g = np.mean(self.reference_image[:,:,1])
                avg_r = np.mean(self.reference_image[:,:,2])
                print(f"Reference avg colors - R:{avg_r:.1f} G:{avg_g:.1f} B:{avg_b:.1f}")
                return True
        except Exception as e:
            print(f"Error loading reference: {e}")
        return False
    
    def match_reference_color(self, current_frame):
        """Calculate white balance to match reference image colors"""
        if self.reference_image is None or current_frame is None:
            print("Need both reference and current frame")
            return
        
        # Use median instead of mean for more robust color matching
        ref_b = np.median(self.reference_image[:,:,0])
        ref_g = np.median(self.reference_image[:,:,1])
        ref_r = np.median(self.reference_image[:,:,2])
        
        # Get current frame WITHOUT white balance applied first
        cur_b = np.median(current_frame[:,:,0])
        cur_g = np.median(current_frame[:,:,1])
        cur_r = np.median(current_frame[:,:,2])
        
        print(f"Current colors - R:{cur_r:.1f} G:{cur_g:.1f} B:{cur_b:.1f}")
        print(f"Reference colors - R:{ref_r:.1f} G:{ref_g:.1f} B:{ref_b:.1f}")
        
        # Calculate ratios needed
        self.sw_blue = (ref_b / cur_b) if cur_b > 0 else 1.0
        self.sw_green = (ref_g / cur_g) if cur_g > 0 else 1.0
        self.sw_red = (ref_r / cur_r) if cur_r > 0 else 1.0
        
        # Clamp values
        self.sw_blue = np.clip(self.sw_blue, 0.3, 3.0)
        self.sw_green = np.clip(self.sw_green, 0.3, 3.0)
        self.sw_red = np.clip(self.sw_red, 0.3, 3.0)
        
        print(f"Calculated WB to match reference:")
        print(f"  R:{self.sw_red:.2f} G:{self.sw_green:.2f} B:{self.sw_blue:.2f}")
    
    def set_hw_white_balance(self, red_ratio, blue_ratio):
        try:
            self.camera.set_string("BalanceRatioSelector", "Red")
            self.camera.set_float("BalanceRatio", red_ratio)
            self.camera.set_string("BalanceRatioSelector", "Blue")
            self.camera.set_float("BalanceRatio", blue_ratio)
            return True
        except:
            return False
    
    def apply_sw_white_balance(self, img):
        """Apply software white balance correction"""
        balanced = img.copy().astype(np.float32)
        balanced[:,:,0] *= self.sw_blue
        balanced[:,:,1] *= self.sw_green
        balanced[:,:,2] *= self.sw_red
        return np.clip(balanced, 0, 255).astype(np.uint8)
    
    def set_exposure(self, exposure_us):
        try:
            self.camera.set_exposure_time(float(exposure_us))
            return True
        except:
            return False
    
    def set_gain(self, gain_db):
        try:
            self.camera.set_gain(float(gain_db))
            return True
        except:
            return False
    
    def capture_frame(self):
        try:
            self.camera.set_acquisition_mode(Aravis.AcquisitionMode.SINGLE_FRAME)
            
            payload = self.camera.get_payload()
            buffer = Aravis.Buffer.new_allocate(payload)
            self.stream.push_buffer(buffer)
            
            self.camera.start_acquisition()
            buffer = self.stream.pop_buffer()
            self.camera.stop_acquisition()
            
            if buffer:
                status = buffer.get_status()
                if status == Aravis.BufferStatus.SUCCESS:
                    data = buffer.get_data()
                    if data and len(data) > 0:
                        img_raw = np.frombuffer(data, dtype=np.uint8).reshape(self.height, self.width)
                        img_color = cv2.cvtColor(img_raw, cv2.COLOR_BAYER_GB2BGR)
                        img_color = self.apply_sw_white_balance(img_color)
                        return img_color
            
            return None
        except Exception as e:
            return None
    
    def analyze_coffee_color(self, img, roi=None):
        if img is None:
            return None
        if roi:
            x, y, w, h = roi
            img = img[y:y+h, x:x+w]
        
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        
        return {
            'hue': np.mean(hsv[:,:,0]),
            'saturation': np.mean(hsv[:,:,1]),
            'value': np.mean(hsv[:,:,2]),
            'lightness': np.mean(lab[:,:,0]),
            'a_channel': np.mean(lab[:,:,1]),
            'b_channel': np.mean(lab[:,:,2])
        }

if __name__ == "__main__":
    print("Initializing camera...")
    cam = InspectionCamera()
    
    if len(sys.argv) > 1:
        cam.load_reference_image(sys.argv[1])
    
    print("\n=== Inspection Camera Ready ===")
    print("\nControls:")
    print("  q - quit")
    print("  s - save image")
    print("  p - save current settings as profile")
    print("  l - load reference image")
    print("  m - match to reference")
    print("  c - analyze coffee color")
    print("  r/R - decrease/increase RED")
    print("  g/G - decrease/increase GREEN")
    print("  b/B - decrease/increase BLUE")
    print("  w - auto white balance (gray world)")
    print("  0 - reset to neutral (1.0, 1.0, 1.0)")
    print("  + - increase exposure")
    print("  - - decrease exposure\n")
    
    exposure = 50000.0
    gain = 1.0
    frame_count = 0
    
    cv2.namedWindow('Camera', cv2.WINDOW_NORMAL)
    cv2.resizeWindow('Camera', 1024, 768)
    
    if cam.reference_image is not None:
        cv2.namedWindow('Reference', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('Reference', 512, 384)
        cv2.imshow('Reference', cv2.resize(cam.reference_image, (512, 384)))
    
    last_frame = None
    
    try:
        while True:
            frame = cam.capture_frame()
            
            if frame is not None:
                frame_count += 1
                last_frame = frame
                
                display = cv2.resize(frame, (1024, 768))
                
                info1 = f"Frame:{frame_count} Exp:{exposure/1000:.1f}ms"
                info2 = f"WB R:{cam.sw_red:.2f} G:{cam.sw_green:.2f} B:{cam.sw_blue:.2f}"
                cv2.putText(display, info1, (10, 30), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                cv2.putText(display, info2, (10, 65), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                
                cv2.imshow('Camera', display)
            
            key = cv2.waitKey(100) & 0xFF
            
            if key == ord('q'):
                break
            elif key == ord('s'):
                if last_frame is not None:
                    filename = f'capture_{int(time.time())}.jpg'
                    cv2.imwrite(filename, last_frame)
                    print(f"✓ Saved {filename}")
            elif key == ord('p'):
                cam.save_profile()
            elif key == ord('l'):
                print("Enter reference image path:")
                filename = input().strip()
                if cam.load_reference_image(filename):
                    cv2.namedWindow('Reference', cv2.WINDOW_NORMAL)
                    cv2.imshow('Reference', cv2.resize(cam.reference_image, (512, 384)))
            elif key == ord('m'):
                if last_frame is not None:
                    cam.match_reference_color(last_frame)
            elif key == ord('c'):
                if last_frame is not None:
                    analysis = cam.analyze_coffee_color(last_frame)
                    if analysis:
                        print(f"\n=== Color Analysis ===")
                        for k, v in analysis.items():
                            print(f"  {k}: {v:.2f}")
            elif key == ord('0'):
                cam.sw_red = cam.sw_green = cam.sw_blue = 1.0
                print("Reset to neutral WB")
            elif key == ord('w'):
                if last_frame is not None:
                    avg_b = np.mean(last_frame[:,:,0])
                    avg_g = np.mean(last_frame[:,:,1])
                    avg_r = np.mean(last_frame[:,:,2])
                    avg_gray = (avg_b + avg_g + avg_r) / 3
                    
                    cam.sw_blue = avg_gray / avg_b if avg_b > 0 else 1
                    cam.sw_green = avg_gray / avg_g if avg_g > 0 else 1
                    cam.sw_red = avg_gray / avg_r if avg_r > 0 else 1
                    print(f"Auto WB: R:{cam.sw_red:.2f} G:{cam.sw_green:.2f} B:{cam.sw_blue:.2f}")
            elif key == ord('r'):
                cam.sw_red = max(0.3, cam.sw_red - 0.05)
                print(f"WB: R:{cam.sw_red:.2f} G:{cam.sw_green:.2f} B:{cam.sw_blue:.2f}")
            elif key == ord('R'):
                cam.sw_red = min(3.0, cam.sw_red + 0.05)
                print(f"WB: R:{cam.sw_red:.2f} G:{cam.sw_green:.2f} B:{cam.sw_blue:.2f}")
            elif key == ord('g'):
                cam.sw_green = max(0.3, cam.sw_green - 0.05)
                print(f"WB: R:{cam.sw_red:.2f} G:{cam.sw_green:.2f} B:{cam.sw_blue:.2f}")
            elif key == ord('G'):
                cam.sw_green = min(3.0, cam.sw_green + 0.05)
                print(f"WB: R:{cam.sw_red:.2f} G:{cam.sw_green:.2f} B:{cam.sw_blue:.2f}")
            elif key == ord('b'):
                cam.sw_blue = max(0.3, cam.sw_blue - 0.05)
                print(f"WB: R:{cam.sw_red:.2f} G:{cam.sw_green:.2f} B:{cam.sw_blue:.2f}")
            elif key == ord('B'):
                cam.sw_blue = min(3.0, cam.sw_blue + 0.05)
                print(f"WB: R:{cam.sw_red:.2f} G:{cam.sw_green:.2f} B:{cam.sw_blue:.2f}")
            elif key == ord('+') or key == ord('='):
                exposure = min(exposure * 1.5, 1000000)
                cam.set_exposure(exposure)
            elif key == ord('-') or key == ord('_'):
                exposure = max(exposure / 1.5, 100)
                cam.set_exposure(exposure)
    
    except KeyboardInterrupt:
        print("\nExiting...")
    finally:
        cv2.destroyAllWindows()

requirements.txt

numpy==2.2.6
opencv-python==4.12.0.88
pillow==12.0.0
pycairo==1.29.0
PyGObject==3.54.5