PRISM2.0 / backend /modules /visual_checks.py
devranx's picture
Initial deploy with LFS images and audio
d790e98
raw
history blame
3.43 kB
import cv2
import numpy as np
from PIL import Image
from backend import config
from backend.utils import get_roi
from backend.model_handler import model_handler
def detect_straight_lines(roi_img):
"""Enhanced edge detection focusing on straight lines."""
gray = cv2.cvtColor(roi_img, cv2.COLOR_BGR2GRAY)
clahe = cv2.createCLAHE(
clipLimit=config.DETECTION_PARAMS['clahe_clip_limit'],
tileGridSize=config.DETECTION_PARAMS['clahe_grid_size']
)
enhanced = clahe.apply(gray)
blurred = cv2.GaussianBlur(
enhanced,
config.DETECTION_PARAMS['gaussian_kernel'],
config.DETECTION_PARAMS['gaussian_sigma']
)
edges = cv2.Canny(
blurred,
config.DETECTION_PARAMS['canny_low'],
config.DETECTION_PARAMS['canny_high']
)
line_mask = np.zeros_like(edges)
lines = cv2.HoughLinesP(
edges,
rho=1,
theta=np.pi/180,
threshold=config.DETECTION_PARAMS['hough_threshold'],
minLineLength=config.DETECTION_PARAMS['min_line_length'],
maxLineGap=config.DETECTION_PARAMS['max_line_gap']
)
if lines is not None:
for line in lines:
x1, y1, x2, y2 = line[0]
cv2.line(line_mask, (x1, y1), (x2, y2), 255, 2)
return line_mask
def simple_edge_detection(roi_img):
"""Simple edge detection."""
gray = cv2.cvtColor(roi_img, cv2.COLOR_BGR2GRAY)
return cv2.Canny(gray, 50, 150)
def ribbon(image_path):
"""Detect the presence of a ribbon in an image."""
image = cv2.imread(image_path)
if image is None:
raise ValueError(f"Could not read image: {image_path}")
h, w = image.shape[:2]
edge_present = []
for i, roi in enumerate(config.ROIS):
x1, y1, x2, y2 = [int(coord * (w if i % 2 == 0 else h)) for i, coord in enumerate(roi)]
roi_img = image[y1:y2, x1:x2]
if i < 6: # Straight line detection for ROIs 0-5
edges = detect_straight_lines(roi_img)
edge_present.append(np.sum(edges) > edges.size * config.DETECTION_PARAMS['edge_pixel_threshold'])
else: # Original method for ROIs 6-8
edges = simple_edge_detection(roi_img)
edge_present.append(np.any(edges))
result = all(edge_present[:6]) and not edge_present[6] and not edge_present[7] and not edge_present[8]
return {"No Ribbon": 0 if result else 1}
def image_quality(image_path):
"""
Check if an image is low resolution or poor quality.
"""
try:
image = Image.open(image_path)
width, height = image.size
pixel_count = width * height
if width < config.MIN_WIDTH or height < config.MIN_HEIGHT or pixel_count < config.MIN_PIXEL_COUNT:
return {"Bad Image Quality": 1}
grayscale_image = image.convert("L")
pixel_array = np.array(grayscale_image)
variance = np.var(pixel_array)
if variance < config.PIXEL_VARIANCE_THRESHOLD:
return {"Bad Image Quality": 1}
return {"Bad Image Quality": 0}
except Exception as e:
print(f"Error processing image: {e}")
return {"Bad Image Quality": 1}
def gnc(image_path):
"""Check for gestures/coach marks and display the image."""
image = get_roi(image_path, *config.GNC)
gnc_text = model_handler.intern(image, config.PGNC, 900).lower()
return {"Visual Gesture or Icon": 1 if 'yes' in gnc_text else 0}